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

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();

A wrong solution:

The problem is similar to the last one, but allows to keep duplicates. So a naive idea is to use Map<Integer, List<Integer>> to store the value and the list of the locations of the value. But why does this solution not work?

For the solution above, we assume that the indices stored in the map are always in an ascending order. So if we swap the value to be deleted, we assume that the list.get(size() - 1) is the tail element of the array. But that's not true. let's take one example:

Insert [1, 1, 2, 2, 3, 3]

Now the hash map is like:

1 -> [0 ,1]

2 -> [2, 3]

3 -> [4, 5]

 

Remove(2)

After remove, the hash map becomes

1->[0, 1]

2->[2]

3 ->[4, 3]

The array is now:

[1, 1, 2, 3, 3]

Then you remove another 2

Remove(2)

So the hash map becomes

1->[0, 1]

2->[]

3->[4, 2]

The array is now

[1, 1, 3, 3]

See the problem? We thought the tail is index 3, however it should be 4. 

So the real problem of using arraylist to store the indices is because we cannot keep the indices in ascending order.

The correct solution:

So the correct solution is to use a HashSet or LinkedHashSet.

Code (Java):

class RandomizedCollection {
    private List<Integer> list;
    private Map<Integer, LinkedHashSet<Integer>> map;
    
    /** Initialize your data structure here. */
    public RandomizedCollection() {
        list = new ArrayList<>();
        map = new HashMap<>();
    }
    
    /** Inserts a value to the collection. Returns true if the collection did not already contain the specified element. */
    public boolean insert(int val) {
        boolean ans = true;
        
        LinkedHashSet<Integer> indices;
        int loc = list.size();
        
        list.add(val);
        
        if (map.containsKey(val)) {
            ans = false;
            indices = map.get(val);
        } else {
            indices = new LinkedHashSet<>();
        }
        indices.add(loc);
        map.put(val, indices);
        
        return ans;
    }
    
    /** Removes a value from the collection. Returns true if the collection contained the specified element. */
    public boolean remove(int val) {
        if (!map.containsKey(val) || map.get(val).isEmpty()) {
            return false;
        }
        
        // Get loc of the val to be removed
        //
        int locToRemove = map.get(val).iterator().next();
        
        // Remove the val
        //
        map.get(val).remove(locToRemove);
        
        if (locToRemove < list.size() - 1) {
        
            // Get the number to be swapped
            //
            int numToSwap = list.get(list.size() - 1);

            // Put the tail number to the location to be removed
            //
            list.set(locToRemove, numToSwap);

            // Update the loc
            //
            if (map.get(numToSwap).contains(list.size() - 1)) {
                map.get(numToSwap).remove(list.size() - 1);
            }
            map.get(numToSwap).add(locToRemove);
        }
        
        // Remove the val
        //
        list.remove(list.size() - 1);
        
        return true;
    }
    
    /** Get a random element from the collection. */
    public int getRandom() {
        if (list.isEmpty()) {
            return 0;
        }
        
        Random rand = new Random();
        int loc = rand.nextInt(list.size());
        return list.get(loc);
    }
}

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