Design a data structure that supports all following operations in average O(1) time.

Note: Duplicate elements are allowed.
insert(val): Inserts an item val to the collection.

remove(val): Removes an item val from the collection if present.

getRandom: Returns a random element from current collection of elements. The probability of each element being returned is linearly related to the number of same value the collection contains.

Example:

// Init an empty collection.
RandomizedCollection collection = new RandomizedCollection();
 
// Inserts 1 to the collection. Returns true as the collection did not contain 1.
collection.insert(1);
 
// Inserts another 1 to the collection. Returns false as the collection contained 1. Collection now contains [1,1].
collection.insert(1);
 
// Inserts 2 to the collection, returns true. Collection now contains [1,1,2].
collection.insert(2);
 
// getRandom should return 1 with the probability 2/3, and returns 2 with the probability 1/3.
collection.getRandom();
 
// Removes 1 from the collection, returns true. Collection now contains [1,2].
collection.remove(1);
 
// getRandom should return 1 and 2 both equally likely.
collection.getRandom();

[Problem Link]

The idea here is to use a HashMap and Array. As we insert val’s, we have a mapping of val’s and nodes (containing value and index). We also insert the node into an array. At any point of time if we need a random value, we simply generate a random number (between 1 to length of array) and return the value from the array (in O(1) time).

How do we remove an element in O(1) time? First we use the map, to find the node where this element is stored. We then swap that with the last element and finally remove the last element. This works because when we look for a random number, we don’t care about order of elements in the array.

This wont work if you don’t store index’es in Nodes.

Let’s look at an example for remove. Below is the status of the HashMap and array at some point of the program. Note that, #xxx is the reference to the Node object (where we store index and value).

1 -> [#abc  #def  #ghi]
2 -> [#jkl  #mno]
3 -> [#pqr]
 
[#abc  #def  #mno  #ghi  #jkl  #pqr]
 
#abc = { idx: 0, val: 1 }
#def = { idx: 1, val: 1 }
#mno = { idx: 2, val: 2 }
#ghi = { idx: 3, val: 1 }
#jkl = { idx: 4, val: 2 }
#pqr = { idx: 5, val: 3 }
 
After remove(2)
 
1 -> [#abc  #def  #ghi]
2 -> [#jkl]
3 -> [#pqr]
 
[#abc  #def  #pqr  #ghi  #jkl]
 
#abc = { idx: 0, val: 1 }
#def = { idx: 1, val: 1 }
#mno = { idx: 2, val: 2 } (DELETED)
#ghi = { idx: 3, val: 1 }
#jkl = { idx: 4, val: 2 }
#pqr = { idx: 2, val: 3 } (UPDATED INDEX AND POSITION IN ARRAY)

Here is what we need todo to remove(2)

  1. Using the HashMap find positions where ‘2’ is stored. We could pick any position to remove it from (we pick #mno pointing to { idx: 2, val: 2 }).
  2. We then pop the last element (i.e. #pqr pointing to { idx: 5, val: 3 })
  3. We need to replace the reference in the array from #mno to #pqr.
    this.arr[removeNode.idx] = lastNode;
  4. Note, at this point #pqr index is still 5, we simply update that to the correct index
    lastNode.idx = removeNode.idx;
/**
 * Initialize your data structure here.
 */
var RandomizedCollection = function() {
    this.arr = [];
    this.map = {};
    this.Node = function(val, idx) {
        this.val = val;
        this.idx = idx;
    }
};
 
/**
 * Inserts a value to the collection. Returns true if the collection did not already contain the specified element. 
 * @param {number} val
 * @return {boolean}
 */
RandomizedCollection.prototype.insert = function(val) {
    if (!this.map[val]) this.map[val] = [];
    let newNode = new this.Node(val, this.arr.length);
    this.arr.push(newNode);
    this.map[val].push(newNode);
    //console.log('insert', val, this.arr);
    return this.map[val].length === 1;
};
 
/**
 * Removes a value from the collection. Returns true if the collection contained the specified element. 
 * @param {number} val
 * @return {boolean}
 */
RandomizedCollection.prototype.remove = function(val) {
    if (!this.map[val] || this.map[val].length === 0) return false;
    let removeNode = this.map[val].pop();
    let lastNode = this.arr.pop();
    if ( removeNode.idx !== lastNode.idx ) {
        lastNode.idx = removeNode.idx;
        this.arr[removeNode.idx] = lastNode;
    }
    //console.log('remove', val, this.map);
    //console.log(this.arr)
    return true;
};
 
/**
 * Get a random element from the collection.
 * @return {number}
 */
RandomizedCollection.prototype.getRandom = function() {
    if (this.arr.length < 1) return;
    let randomIdx = Math.floor(Math.random() * this.arr.length);
    //console.log('random', randomIdx, this.arr, this.arr.length)
    return this.arr[randomIdx].val;
};
 
/** 
 * Your RandomizedCollection object will be instantiated and called as such:
 * var obj = new RandomizedCollection()
 * var param_1 = obj.insert(val)
 * var param_2 = obj.remove(val)
 * var param_3 = obj.getRandom()
 */