insert-delete-getrandom-o1-duplicates-allowed.java

// Problem link - https://leetcode.com/problems/insert-delete-getrandom-o1-duplicates-allowed/


class RandomizedCollection {
    List<Integer> nums;
    Map<Integer, Set<Integer>> idxMap;
    Random random;

    public RandomizedCollection() {
        nums = new ArrayList<>();
        idxMap = new HashMap<>();
        random = new Random();    
    }
    
    public boolean insert(int val) {
        boolean doesntContain = !idxMap.containsKey(val);

        if(doesntContain){
            idxMap.put(val, new HashSet<>());
        } 

        idxMap.get(val).add(nums.size());
        nums.add(val);

        return doesntContain;
    }
    
    public boolean remove(int val) {
        if(!idxMap.containsKey(val)){
            return false;
        }

        Set<Integer> idxSet = idxMap.get(val);
        int idxToBeRemoved = idxSet.iterator().next();
        if(idxSet.size() == 1){
            idxMap.remove(val);
        } else {
            idxSet.remove(idxToBeRemoved);
        }

        int lastIdx = nums.size() - 1;
        if(idxToBeRemoved != lastIdx){
            int lastVal = nums.get(lastIdx);
            Set<Integer> lastIdxSet = idxMap.get(lastVal);
            lastIdxSet.add(idxToBeRemoved);
            lastIdxSet.remove(lastIdx);
            nums.set(idxToBeRemoved, lastVal);
        }

        nums.remove(lastIdx);

        return true;
    }
    
    public int getRandom() {
        return nums.get(random.nextInt(nums.size()));
    }
}

/**
 * Your RandomizedCollection object will be instantiated and called as such:
 * RandomizedCollection obj = new RandomizedCollection();
 * boolean param_1 = obj.insert(val);
 * boolean param_2 = obj.remove(val);
 * int param_3 = obj.getRandom();
 */