门 | GATE-CS-2003 | 问题 7

本题目要求实现一种数据结构，该数据结构具有如下操作：

• insert(x): 如果数据结构中没有值为 x 的元素，则将 x 插入到数据结构中。
• delete(x): 如果数据结构中存在值为 x 的元素，则将其从数据结构中删除。
• getRandom(): 随机返回数据结构中的一个元素，每个元素被返回的概率相等。

要求实现的数据结构应该支持上述操作，且时间复杂度为 $O(1)$。

解决方案

在保持时间复杂度为 $O(1)$ 的前提下，我们需要一些概率论上的技巧来实现本题目的要求。我们可以利用哈希表和数组来实现如下的数据结构。

• 哈希表：用来记录每个元素在数组中的下标（如果数组中不存在该元素，则在哈希表中添加一个键值对）。

• 数组：用来保存元素。我们需要有一个变量 size 来记录数组中的有效元素个数。

操作 insert(x)：

• 如果 x 在哈希表中已经存在，则返回。
• 在数组的末尾插入 x，并在哈希表中添加一个键值对 (x, size)，其中 size 是插入前数组的元素总数。
• 将 size 加 1。

操作 delete(x)：

• 如果 x 在哈希表中不存在，则返回。
• 将数组中的 x 替换为最后一个元素 y，并更新哈希表中的键值对 (y, size')，其中 size' 是删除 x 后的数组元素个数。
• 在哈希表中删除键值对 (x, size)。
• 将 size 减 1。

操作 getRandom()：

• 随机生成一个整数 i，其中 $0 \leq i < size$。
• 返回数组中下标为 i 的元素。

示例代码

import random

class RandomizedSet:
    def __init__(self):
        self.nums = []
        self.index = {}

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

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

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

上述代码使用 Python 实现了一个基于数组和哈希表的随机集合数据结构。其中，nums 保存了集合中的元素，index 是一个哈希表，记录了元素在 nums 中出现的下标。具体实现见代码注释。

注意，在 Python 中，append 和 pop 操作的时间复杂度均为 $O(1)$，因此本题目中的所有操作都可以在 $O(1)$ 时间内完成，符合题目要求。