Leetcode 381: Insert Delete GetRandom O(1) - Duplicates allowed

https://leetcode.com/problems/insert-delete-getrandom-o1-duplicates-allowed/description/

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

Note: Duplicate elements are allowed.

1. insert(val): Inserts an item val to the collection.
2. remove(val): Removes an item val from the collection if present.
3. 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();

Notes:

This question can be viewed as a follow up to Leetcode 380: Insert Delete GetRandom O(1).

If we still use a vector to keep all the elements, the getRandom() should be still in O(1). But now we need to deal with the duplicates.

One identifier to duplicates are their indexes in the vector. So one way is that we can use the index as one part of the key of the hashmap.

See the code below:

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) {
        int len = v.size();
        ump[val][len] = 1;
        v.push_back(val);
        return ump[val].size() == 1;
    }

    /** Removes a value from the collection. Returns true if the collection contained the specified element. */
    bool remove(int val) {
        if(!ump.count(val) || ump[val].empty() ) return false;
        int id = ump[val].begin()->first, a = v.back(), len = v.size();
        v[id] = a;
        ump[val].erase(id);
        ump[a][id] = 1;
        ump[a].erase(len-1);
        v.pop_back();
        return true;
    }

    /** Get a random element from the collection. */
    int getRandom() {
        return v[rand() % v.size()];
    }

private:
    unordered_map<int, unordered_map<int, int>> ump;//<val, <indx, ct>>
    vector<int> v;
};

/**
 * Your RandomizedCollection object will be instantiated and called as such:
 * RandomizedCollection* obj = new RandomizedCollection();
 * bool param_1 = obj->insert(val);
 * bool param_2 = obj->remove(val);
 * int param_3 = obj->getRandom();
 */

Another way is that we can define the following data structure:

unordered_map<int, vector<int>> mp;  // <val, <index>>;

Then how can we localize the index of the element inside <index> above? We can store this information inside the vector:

vector<pair<int , int>> //<pair<val, index inside mp[val] or <index>>

When a new element is inserted:

mp[val].push_back(v.size());
v.push_back({val, mp[val].size() - 1});

In this way, mp[val[i].first][val[i].second] = i. Thus, from the element inside the v, we can localize the element position inside the mp <int, vector<int>>.