[Leetcode] Problem 381 - Insert Delete GetRandom O(1) - Duplicates allowed

Posted on | In ,

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

  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.

Note

Duplicate elements are allowed.

Example

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
// 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();

Code

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
private Random random;
private Map<Integer, List<Integer>> map;
private List<int[]> vals;

public RandomizedCollection() {
    random = new Random();
    map = new HashMap<>();
    vals = new ArrayList<>();
}

public boolean insert(int val) {
    map.computeIfAbsent(val, l -> new ArrayList<>()).add(vals.size());
    vals.add(new int[] {val, map.get(val).size() - 1});
    return map.get(val).size() == 1;
}

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

    int[] lastVal = vals.get(vals.size() - 1);
    List<Integer> index = map.get(val);
    int idx = index.get(index.size() - 1);

    vals.set(idx, lastVal);
    map.get(lastVal[0]).set(lastVal[1], idx);
    index.remove(index.size() - 1);
    vals.remove(vals.size() - 1);

    if (index.isEmpty())
        map.remove(val);

    return true;
}

public int getRandom() {
    int idx = random.nextInt(vals.size());
    return vals.get(idx)[0];
}