381. Insert Delete GetRandom O(1) - Duplicates allowed

class RandomizedCollection {

    private final Map<Integer, Set<Integer>> valueToIndexSetMap;
    private final List<Integer> valueList;

    public RandomizedCollection() {
        this.valueToIndexSetMap = new HashMap<>();
        this.valueList = new ArrayList<>();
    }

    public boolean insert(int val) {
        boolean exist = valueToIndexSetMap.containsKey(val);
        valueToIndexSetMap.computeIfAbsent(val, key -> new HashSet<>()).add(valueList.size());
        valueList.add(val);
        return !exist;
    }

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

        Set<Integer> indexSet = valueToIndexSetMap.get(val);
        Iterator<Integer> indexSetIterator = indexSet.iterator();
        int removedIndex = indexSetIterator.next();
        indexSetIterator.remove();
        if (indexSet.isEmpty()) {
            valueToIndexSetMap.remove(val);
        }

        // 注意仅当删除的元素不是最后一个元素时才需要将最后一个元素进行交换处理
        if (removedIndex != valueList.size() - 1) {
            // 注意因为 lastValue 挪动了位置，所以需要同步修改 valueToIndexListMap 中的索引
            int lastValue = valueList.get(valueList.size() - 1);
            valueList.set(removedIndex, lastValue);

            Set<Integer> lastValueIndexSet = valueToIndexSetMap.get(lastValue);
            lastValueIndexSet.remove(valueList.size() - 1);
            if (lastValueIndexSet.isEmpty()) {
                valueToIndexSetMap.remove(lastValue);
            }
            valueToIndexSetMap.computeIfAbsent(lastValue, key -> new HashSet<>()).add(removedIndex);
        }

        valueList.remove(valueList.size() - 1);
        return true;
    }

    public int getRandom() {
        return valueList.get(ThreadLocalRandom.current().nextInt(valueList.size()));
    }

}

References

381. Insert Delete GetRandom O(1) - Duplicates allowed