Insert Delete GetRandom O(1)

Description

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.

No duplicates

The idea is to store the value itself in a vector(without duplicates of course), and then use a hash map to save (val, the index of val in the vector). The hardest part should be the delete() operation: it should swap the index of the last element and the to-be-deleted element, e.g. (1, 1), (2, 2), (3, 3), we delete the (2, 2), and then we have (1, 1), (3, 2)

  private:
        vector<int> nums;
        unordered_map<int, int> m;
  /*
  * Removes a value from the set. Returns true if the set 
  * contained the specified element. 
  */
    bool remove(int val) {
        if (m.find(val) == m.end()) return false;
        int last = nums.back();
        m[last] = m[val];
        nums[m[val]] = last;
        nums.pop_back();
        m.erase(val);
        return true;
    }

With duplicates

Honestly, kind of hard to visualize the delete operation at first. We need a bucket(which is implemented by hash map) to store (value, position in the nums array) and a vector to store (value, position in the value bucket). For example, we insert 1, 1, 2, 3 sequentially into this data structure, and we would have:

• buf[1] : ${0, 1}$, buf[2] : ${2}$, buf[3] : ${3}$
• nums: $$(1, 0), (1, 1), (2, 0), (3, 0)$$

Then the delete(1) operation would do something similar to the previous question:
1. get nums.back(), i.e. (3, 0) as variable last
2. buf[3][0]would be replaced by buf[1].back(), i.e. 1., and it would become buf[3] :{1}
3. nums[1] = last, (1, 1) = (3, 0), and nums becomes: $(1, 0), (1, 1), (2, 0), (3, 0)$
4. buf[1].pop_back() and nums.pop_back()

class RandomizedCollection {
public:
    /** Initialize your data structure here. */
    RandomizedCollection() {     
    }

    /** Inserts a value to the collection. Returns true if the collection did not already contain the specified element. */
    bool insert(int val) {
        buf[val].push_back(nums.size());
        nums.push_back(make_pair(val, buf[val].size() - 1));
        if(!buf.count(val))  return true;
        return false;
    }

    /** Removes a value from the collection. Returns true if the collection contained the specified element. */
    bool remove(int val) {
        if(buf.count(val)){
            auto last = nums.back();
            buf[last.first][last.second] = buf[val].back();
            nums[buf[val].back()] = last;
            buf[val].pop_back();
            if(buf[val].empty())
                buf.erase(val);
            nums.pop_back();
            return true;
        }
        return false;
    }


    /** Get a random element from the collection. */
    int getRandom() {
        return nums[rand() % nums.size()].first;
    }
private:
    unordered_map<int, vector<int>> buf;
    vector<pair<int, int>> nums;
};