Insert Delete GetRandom O(1) - Duplicates allowed LeetCode solution in Java

1
class RandomizedCollection {
2
  List<Integer> multiSet;
3
  HashMap<Integer, PriorityQueue<Integer>> map;
4
  Random random;
5

6
  public RandomizedCollection() {
7
    map = new HashMap<>();
8
    multiSet = new ArrayList<>();
9
    random = new Random();
10
  }
11

12
  public boolean insert(int val) {
13
    boolean contains = map.containsKey(val);
14
    multiSet.add(val);
15
    PriorityQueue<Integer> pq;
16
    if (!contains) {
17
      pq = new PriorityQueue<>(Collections.reverseOrder());
18
      map.put(val, pq);
19
    } else pq = map.get(val);
20

21
    pq.add(multiSet.size() - 1);
22
    return !contains;
23
  }
24

25
  public boolean remove(int val) {
26
    if (!map.containsKey(val)) return false;
27

28
    PriorityQueue<Integer> pq = map.get(val);
29
    int indexToRemove = pq.poll();
30
    if (pq.size() == 0) map.remove(val);
31

32
    int lastIndex = multiSet.size() - 1;
33
    if (indexToRemove != lastIndex) {
34
      int valLast = multiSet.get(lastIndex);
35
      PriorityQueue<Integer> temp = map.get(valLast);
36
      temp.poll();
37
      temp.add(indexToRemove);
38
      multiSet.set(indexToRemove, valLast);
39
    }
40
    multiSet.remove(lastIndex);
41
    return true;
42
  }
43

44
  public int getRandom() {
45
    int index = random.nextInt(multiSet.size());
46
    return multiSet.get(index);
47
  }
48
}