LeetCode – Insert Delete GetRandom O(1) (Java)

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

insert(val): Inserts an item val to the set if not already present.
remove(val): Removes an item val from the set if present.
getRandom: Returns a random element from current set of elements. Each element must have the same probability of being returned.

Java Solution

We can use two hashmaps to solve this problem. One uses value as keys and the other uses index as the keys.

class RandomizedSet {
    HashMap<Integer, Integer> valueMap;
    HashMap<Integer, Integer> idxMap;
 
    /** Initialize your data structure here. */
    public RandomizedSet() {
        valueMap = new HashMap<>();
        idxMap = new HashMap<>();
    }
 
    /** Inserts a value to the set. Returns true if the set did not already contain the specified element. */
    public boolean insert(int val) {
        if(valueMap.containsKey(val)){
            return false;
        }
 
        valueMap.put(val, valueMap.size());
        idxMap.put(idxMap.size(), val);
 
        return true;
    }
 
    /** Removes a value from the set. Returns true if the set contained the specified element. */
    public boolean remove(int val) {
        if(valueMap.containsKey(val)){
            int idx = valueMap.get(val);
            valueMap.remove(val);
            idxMap.remove(idx);
 
            Integer tailElem = idxMap.get(idxMap.size());
            if(tailElem!=null){
                idxMap.put(idx,tailElem);
                valueMap.put(tailElem, idx);
            }
 
            return true;
        }
 
        return false;
    }
 
    /** Get a random element from the set. */
    public int getRandom() {
        if(valueMap.size()==0){
            return -1;
        }
 
        if(valueMap.size()==1){
            return idxMap.get(0);
        }
 
        Random r = new Random();
        int idx = r.nextInt(valueMap.size());
 
        return idxMap.get(idx);
    }
}
Category >> Algorithms  
If you want someone to read your code, please put the code inside <pre><code> and </code></pre> tags. For example:
<pre><code> 
String foo = "bar";
</code></pre>
  • Hooman

    OR simply use LinkedHashMap which can return a number by index

  • Cherry Zhao

    Thanks for sharing the code. This question is interesting and has been asked by Uber recently. I also read this in-depth analysis of the same question here:

    http://bit.ly/2bCPP47