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

标签: 设计 数组 哈希表 数学 随机化

难度: Medium

实现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 方法时，数据结构中 至少存在一个 元素。

Submission

运行时间: 247 ms

内存: 0.0 MB

import random

class RandomizedSet:

    def __init__(self):
        self.val_to_index = {}
        self.values = []

    def insert(self, val: int) -> bool:
        if val in self.val_to_index:
            return False
        self.values.append(val)
        self.val_to_index[val] = len(self.values) - 1
        return True

    def remove(self, val: int) -> bool:
        if val not in self.val_to_index:
            return False
        last_value, index = self.values[-1], self.val_to_index[val]
        self.values[index], self.val_to_index[last_value] = last_value, index
        self.values.pop()
        del self.val_to_index[val]
        return True

    def getRandom(self) -> int:
        return random.choice(self.values)

# Your RandomizedSet object will be instantiated and called as such:
# obj = RandomizedSet()
# param_1 = obj.insert(val)
# param_2 = obj.remove(val)
# param_3 = obj.getRandom()

Explain

这个题解使用了一个列表 values 和一个字典 val_to_index 来实现 O(1) 时间复杂度的插入、删除和随机获取元素。values 列表按插入顺序存储元素，val_to_index 字典则记录每个元素在 values 列表中的索引。插入元素时，将其添加到 values 列表末尾，并在 val_to_index 中记录其索引。删除元素时，先将要删除的元素与 values 列表的最后一个元素交换位置，然后删除列表最后一个元素并更新 val_to_index。获取随机元素时，直接在 values 列表中随机选择一个元素返回。

时间复杂度: O(1)

空间复杂度: O(n)

import random

class RandomizedSet:

    def __init__(self):
        # 用于存储元素值到索引的映射
        self.val_to_index = {}
        # 用于按插入顺序存储元素
        self.values = []

    def insert(self, val: int) -> bool:
        # 如果元素已存在，返回 False
        if val in self.val_to_index:
            return False
        # 将元素添加到列表末尾
        self.values.append(val)
        # 在字典中记录元素的索引
        self.val_to_index[val] = len(self.values) - 1
        return True

    def remove(self, val: int) -> bool:
        # 如果元素不存在，返回 False
        if val not in self.val_to_index:
            return False
        # 获取要删除元素的索引和列表最后一个元素
        last_value, index = self.values[-1], self.val_to_index[val]
        # 将要删除的元素与最后一个元素交换位置
        self.values[index], self.val_to_index[last_value] = last_value, index
        # 删除列表的最后一个元素
        self.values.pop()
        # 从字典中删除元素
        del self.val_to_index[val]
        return True

    def getRandom(self) -> int:
        # 从列表中随机选择一个元素返回
        return random.choice(self.values)

Explore

在`remove`操作中，将要删除的元素与列表的最后一个元素交换位置是为了保持列表的连续性，从而可以在O(1)的时间复杂度内完成删除操作。如果不进行交换，直接删除中间的元素会导致列表中间出现空位，或者需要移动该元素之后的所有元素来填补空位，这将使删除操作的时间复杂度上升到O(n)。通过交换，我们可以确保删除的总是列表的最后一个元素，这样只需要更新一次列表和字典，极大地提高了效率。

是的，随机获取元素时，所有元素的选择概率确实都是相同的。这是通过在列表中均匀随机选择一个索引来实现的。由于列表`values`中的每个元素都正好占据一个位置，且这些位置在插入时是连续分配的，因此使用`random.choice(self.values)`函数可以保证从这个列表中随机选出的每个元素的概率都是相等的，即每个元素被选择的概率都是1/N，其中N是列表中元素的总数。

如果删除的元素恰好是列表中的最后一个元素，那么不需要进行交换操作。这种情况下的处理逻辑非常简单：直接移除列表中的最后一个元素，并从字典`val_to_index`中删除该元素的索引即可。因为没有其他元素需要交换位置，这种情况下的删除操作依然可以在O(1)时间复杂度内完成。

在删除操作中，交换元素后只需更新最后一个元素的索引，而不需要更新被删除元素的索引，因为被删除元素的索引信息将从字典中移除。在交换过程中，原本位于列表末尾的元素移至被删除元素的位置，因此我们只需要更新这个元素在字典中的索引。被删除的元素由于已经不再列表中，其索引信息不再需要，因此直接从字典中删除即可。这样的处理保持了字典的准确性和操作的高效性。