🌕🌗 381. O(1) 时间插入、删除和获取随机元素 - 允许重复

难度: 🌕🌗

问题描述

---

解法

class RandomizedCollection {
    // 思路：
    // HashMap 存储 val - Set(index)
    // ArrayList 存储 val
    HashMap<Integer, HashSet<Integer>> map;
    ArrayList<Integer> list;

    public RandomizedCollection() {
        map = new HashMap<>();
        list = new ArrayList<>();
    }
    
    public boolean insert(int val) {
        boolean res = false;
        if(!map.containsKey(val)) {
            HashSet<Integer> tmp = new HashSet<>();
            map.put(val, tmp);
            res = true;
        }
        HashSet<Integer> tmpList = map.get(val);
        list.add(val);
        int size = list.size();
        tmpList.add(size - 1);
        map.put(val, tmpList);
        return res;
    }
    
    public boolean remove(int val) {
        if(!map.containsKey(val)) {
            return false;
        }
        HashSet<Integer> tmpList = map.get(val);
        int size = tmpList.size();
        // 通过 Iterator 获取 set 中一个下标
        Iterator<Integer> it = tmpList.iterator();
        int rvIndex = it.next(); 
        tmpList.remove(rvIndex);
        if(size == 1) {
            map.remove(val);
        } else {
            map.put(val, tmpList);
        }

        // 针对 list 的操作
        int len = list.size();
        if(rvIndex == len - 1) {
            // 删除的正好是总链表的最后一个元素
            list.remove(len - 1);
            return true;
        }
        int lastVal = list.get(len - 1); // 将 list 最后一个元素的值替换到 list 对应位置，那么 map 相关也需要改变
        list.set(rvIndex, lastVal);
        list.remove(len - 1);
        tmpList = map.get(lastVal);
        // 删除 set 中对应 list 的最后一个下标，替换成当前新下标即可
        tmpList.remove(len - 1);
        tmpList.add(rvIndex);
        map.put(lastVal, tmpList);
        return true;
    }
    
    public int getRandom() {
        int len = list.size();
        int index = new Random().nextInt(len);
        return list.get(index);
    }
}

/**
 * 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();
 */

---

输出