Medium | LeetCode 380. 常数时间插入、删除和获取随机元素 | 设计

380. 常数时间插入、删除和获取随机元素

设计一个支持在平均 时间复杂度 O(1) 下，执行以下操作的数据结构。

1. insert(val)：当元素 val 不存在时，向集合中插入该项。
2. remove(val)：元素 val 存在时，从集合中移除该项。
3. getRandom：随机返回现有集合中的一项。每个元素应该有相同的概率被返回。

示例 :

// 初始化一个空的集合。
RandomizedSet randomSet = new RandomizedSet();

// 向集合中插入 1 。返回 true 表示 1 被成功地插入。
randomSet.insert(1);

// 返回 false ，表示集合中不存在 2 。
randomSet.remove(2);

// 向集合中插入 2 。返回 true 。集合现在包含 [1,2] 。
randomSet.insert(2);

// getRandom 应随机返回 1 或 2 。
randomSet.getRandom();

// 从集合中移除 1 ，返回 true 。集合现在包含 [2] 。
randomSet.remove(1);

// 2 已在集合中，所以返回 false 。
randomSet.insert(2);

// 由于 2 是集合中唯一的数字，getRandom 总是返回 2 。
randomSet.getRandom();

解题思路

首先分析题目需要使用什么数据结构。

1. O(1) 时间内插入, 数组, 链表, 哈希Map, 均支持O(1)时间的插入

2. O(1) 时间删除, 常规的数组删除元素, 复杂度是O(N), 链表的删除操作复杂度是O(1)。不过都需要事先定位到需要删除的位置。可以使用HashMap保存位置。

3. O(1) 时间复杂度获取随机数。这个使用数组是最合适的。用链表不太好获取到随机位置的数。

综上, 使用数组存储数据相对来讲更加合适。那么如何才能在O(1)时间删除元素呢？很简单, 常规的删除操作是将后边的部分全部向前移动一部分。这个时间复杂度是O(N)。实际上一个更加方便的删除方法是, 直接把数组的最后一个元素拿过来填充在要删除的位置, 然后再把最后一个位置的元素删除掉。当然, 这个时候, 也需要更新用于定位元素的hashmap索引。

class RandomizedSet {

    private final Map<Integer, Integer> index;
    private final List<Integer> data;
    private final Random random;

    /** Initialize your data structure here. */
    public RandomizedSet() {
        index = new HashMap<>();
        data = new ArrayList<>();
        random = new Random();
    }

    /** Inserts a value to the set. Returns true if the set did not already contain the specified element. */
    public boolean insert(int val) {
        if (index.containsKey(val)) {
            return false;
        }
        // 直接将元素添加在数组末尾, 并且添加索引
        index.put(val, data.size());
        data.add(val);
        return true;
    }

    /** Removes a value from the set. Returns true if the set contained the specified element. */
    public boolean remove(int val) {
        Integer idx = index.get(val);
        if (idx != null) {
            int size = data.size();
            // 用数组最后一个元素填充要删除的位置
            data.set(idx, data.get(size - 1));
            // 更新索引
            index.put(data.get(idx), idx);
            // 删除数组最后一个元素
            data.remove(size - 1);
            // 删除被删除元素的索引
            index.remove(val);
            return true;
        }
        return false;
    }

    /** Get a random element from the set. */
    public int getRandom() {
        return data.get(random.nextInt(data.size()));
    }
}