Design a Container with O(1) Add, Remove and GetRandomElement

Task: Design a Data Structure (Container) that supports O(1) constant time in Adding, Removing and Getting a Random Element with Equal Probability. For example:

// RandomDS.add(3)
// RandomDS.add(4)
// RandomDS.add(10)
// RandomDS.add(3)
// RandomDS.getRandomElement(), 3: 50%, 4: 25%, 10: 25%

In order to achieve O(1) in Add, Remove and GetRandomElement, we use two data structures, the first is the vector of the data, which stores all numbers including the duplicates, thus easy to return a random element with equal probabilities (duplicates have higher probabilities of being selected).

The second data structure is a hash map, that maps the number (values) to a collection (set) of indices. Adding is trivial O(1) as we only need to push_back the new element to the vector and insert the index to the set.

Removing is a bit tricky. First, we can get a index for the number, then swapping it out with the last element of the vector. Then, we need to update the indices for the both elements accordingly. Lastly, we need to shrink the vector by one element, as to remove the number from the vector.

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
40
41
42

class RandomDS {
public:  
  void add(int num) {
    data.push_back(num);
    cache[num].insert(data.size() - 1);
  }
  
  bool remove(int num) {
    if ((cache.find(num) == cache.end()) || (cache[num].size() == 0)) return false;
    // get the first index of the num
    int an_index = *cache[num].begin(); 
    swap(data[an_index], data[data.size() - 1]);
    // to update the index of the number that is swapped in.
    int x = data[data.size() - 1];    
    // update the index of the number that is swapped in.
    cache[x].erase(data.size() - 1);
    cache[x].insert(an_index);
    // remove the index of num
    cache[x].erase(an_index);
    // remove the num (last element) from the vector
    data.pop_back();
    return true;
  }
    
  int getRandomElement() {
    if (data.empty()) {
       throw std::exception("empty."); 
    }
    int n = data.size();
    int r = rnd() % n;
    return data[r];
  }
  
  int getSize() {
     return data.size(); 
  }
private:
  // store all the numbers with duplicates
  vector<int> data;  
  // store the indices mapping of all numbers
  unordered_map<int, unordered_set<int>> cache;
}

class RandomDS {
public:  
  void add(int num) {
    data.push_back(num);
    cache[num].insert(data.size() - 1);
  }
  
  bool remove(int num) {
    if ((cache.find(num) == cache.end()) || (cache[num].size() == 0)) return false;
    // get the first index of the num
    int an_index = *cache[num].begin(); 
    swap(data[an_index], data[data.size() - 1]);
    // to update the index of the number that is swapped in.
    int x = data[data.size() - 1];    
    // update the index of the number that is swapped in.
    cache[x].erase(data.size() - 1);
    cache[x].insert(an_index);
    // remove the index of num
    cache[x].erase(an_index);
    // remove the num (last element) from the vector
    data.pop_back();
    return true;
  }
    
  int getRandomElement() {
    if (data.empty()) {
       throw std::exception("empty."); 
    }
    int n = data.size();
    int r = rnd() % n;
    return data[r];
  }
  
  int getSize() {
     return data.size(); 
  }
private:
  // store all the numbers with duplicates
  vector<int> data;  
  // store the indices mapping of all numbers
  unordered_map<int, unordered_set<int>> cache;
}

Unit Test Cases

Below are some test cases if you want to unit test or (TDD Test Driven Development).

• add 1, 2, 3, remove 3, 2, 1 , getSize() should be 0, remove 1 more, exception
• add 1, 2, 3, getRandomElement() – 10000, count frequency of 1, 2, and 3. 3333, with margin, 20
• add 1, 2, 3, 4, do the above, remove 4, getRandomElement(), count the frequency,
• add 10 1’s , add 2 2’s, count the freq,
• add 1, 2 remove 1, 2.

–EOF (The Ultimate Computing & Technology Blog) —