380. O(1) 时间插入、删除和获取随机元素

题目

题目链接

实现RandomizedSet 类：

• RandomizedSet() 初始化 RandomizedSet 对象
• bool insert(int val) 当元素 val 不存在时，向集合中插入该项，并返回 true ；否则，返回 false 。
• bool remove(int val) 当元素 val 存在时，从集合中移除该项，并返回 true ；否则，返回 false 。
• int getRandom() 随机返回现有集合中的一项（测试用例保证调用此方法时集合中至少存在一个元素）。每个元素应该有 相同的概率 被返回。

你必须实现类的所有函数，并满足每个函数的 平均 时间复杂度为 O(1) 。

示例：

输入
["RandomizedSet", "insert", "remove", "insert", "getRandom", "remove", "insert", "getRandom"]
[[], [1], [2], [2], [], [1], [2], []]
输出
[null, true, false, true, 2, true, false, 2]

解释
RandomizedSet randomizedSet = new RandomizedSet();
randomizedSet.insert(1); // 向集合中插入 1 。返回 true 表示 1 被成功地插入。
randomizedSet.remove(2); // 返回 false ，表示集合中不存在 2 。
randomizedSet.insert(2); // 向集合中插入 2 。返回 true 。集合现在包含 [1,2] 。
randomizedSet.getRandom(); // getRandom 应随机返回 1 或 2 。
randomizedSet.remove(1); // 从集合中移除 1 ，返回 true 。集合现在包含 [2] 。
randomizedSet.insert(2); // 2 已在集合中，所以返回 false 。
randomizedSet.getRandom(); // 由于 2 是集合中唯一的数字，getRandom 总是返回 2 。

提示：

• -231 <= val <= 231 - 1
• 最多调用 insert、remove 和 getRandom 函数 2 * 105 次
• 在调用 getRandom 方法时，数据结构中 至少存在一个 元素。

解法

O(1) 时间实现插入和删除借助 Map 或者 Set 就可以很容易的实现，难点是如何 O(1) 时间获取随机项。

O(1) 时间获取随机项必须使用数组这样可以通过下标直接获取对应位置元素。

那么如果使用 Map<val, index>，元素值作为 key，在数组中的位置作为 value，同时使用 ArrayList<Integer> 保存元素值，这样存在一个问题，删除元素的时候 list 中被删除元素之后的元素下标会 -1，而 Map 中还需要对变更的数组下标进行维护，这样显然不能做到 O(1) 时间复杂度。

所以重点是如何避免删除元素后，后续元素下标产生变化。可以将数组最后一位的元素填充到被删除的位置，这样只有最后一位元素的下标产生了变化，其余元素都没有发生变化。

代码如下：

class RandomizedSet {

    Map<Integer, Integer> map;
    List<Integer> list;
    Random random;

    public RandomizedSet() {
        map = new HashMap<>();
        list = new ArrayList<>();
        random = new Random();
    }

    public boolean insert(int val) {
        if (map.containsKey(val)) {
            return false;
        }
        list.add(val);
        map.put(val, list.size() - 1);
        return true;
    }

    public boolean remove(int val) {
        if (!map.containsKey(val)) {
            return false;
        }
        // 如果删除的不是最后一位元素，则先将最后一位元素复制到被删除元素的位置
        int index = map.get(val);
        int lastElementIndex = list.size() - 1;
        if (index < lastElementIndex) {
            int lastElementVal = list.get(lastElementIndex);
            list.set(index, lastElementVal);
            map.put(lastElementVal, index);
        }
        // 删除最后一位元素
        list.remove(lastElementIndex);
        map.remove(val);
        return true;
    }

    public int getRandom() {
        int randomIndex = random.nextInt(list.size());
        return list.get(randomIndex);
    }
}

使用数组的解法：

class RandomizedSet {
    static int[] nums = new int[200010];
    Random random = new Random();
    Map<Integer, Integer> map = new HashMap<>();
    int idx = -1;
    public boolean insert(int val) {
        if (map.containsKey(val)) return false;
        nums[++idx] = val;
        map.put(val, idx);
        return true;
    }
    public boolean remove(int val) {
        if (!map.containsKey(val)) return false;
        int loc = map.remove(val);
        if (loc != idx) map.put(nums[idx], loc);
        nums[loc] = nums[idx--];
        return true;
    }
    public int getRandom() {
        return nums[random.nextInt(idx + 1)];
    }
}