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

by

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);
    }
}

2 thoughts on “LeetCode – Insert Delete GetRandom O(1) (Java)”

Leave a Comment