Insert Delete GetRandom O(1) – Duplicates allowed

leetcode 381.

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.

Solution. Same as previous problem. To solve the duplicate one, all we need to change is that we choose to use HashSet as the value for loc. When we do remove and update, we 1. remove HashSet 2. remove ArrayList 3. Update HashSet 4.Update ArrayList.

public class RandomizedCollection {

    Map<Integer, Set<Integer>> loc;
    List<Integer> list;
    Random random = new Random();

    /** Initialize your data structure here. */
    public RandomizedCollection() {
        loc = new HashMap<>();
        list = new ArrayList<>();
    }

    /** Inserts a value to the collection. Returns true if the collection did not already contain the specified element. */
    public boolean insert(int val) {
        Set<Integer> set = loc.getOrDefault(val, new HashSet<>());
        set.add(list.size());
        loc.put(val, set);
        list.add(val);
        return true;
    }

    /** Removes a value from the collection. Returns true if the collection contained the specified element. */
    public boolean remove(int val) {
        if (!loc.containsKey(val)) {
            return false;
        }
        Set<Integer> set = loc.get(val);
        int pos = set.iterator().next();
        set.remove(pos);    // remove loc
        if (set.isEmpty()) {
            loc.remove(val);
        }
        int lastEle = list.remove(list.size() - 1); // remove list
        if (pos != list.size()) {   // update list and loc
            list.set(pos, lastEle);     // update loc
            Set<Integer> lastEleSet = loc.get(lastEle);
            lastEleSet.remove(list.size()); // update list
            lastEleSet.add(pos);
        }
        return true;
    }

    /** Get a random element from the collection. */
    public int getRandom() {
        return list.get(random.nextInt(list.size()));
    }

}

Check my code on github.