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

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

---

题目来源：力扣（LeetCode）https://leetcode-cn.com/problems/insert-delete-getrandom-o1-duplicates-allowed/

题目

---

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

注意: 允许出现重复元素。

1. insert(val)：向集合中插入元素 val。
2. remove(val)：当 val 存在时，从集合中移除一个 val。
3. getRandom：从现有集合中随机获取一个元素。每个元素被返回的概率应该与其在集合中的数量呈线性相关。

示例:

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

// 向集合中插入 1 。返回 true 表示集合不包含 1 。
collection.insert(1);

// 向集合中插入另一个 1 。返回 false 表示集合包含 1 。集合现在包含 [1,1] 。
collection.insert(1);

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

// getRandom 应当有 2/3 的概率返回 1 ，1/3 的概率返回 2 。
collection.getRandom();

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

// getRandom 应有相同概率返回 1 和 2 。
collection.getRandom();

解题思路

---

思路：哈希表

先审题，现在题目要求设计支持在平均时间复杂度为 O(1) 的情况下 ，执行插入、删除、随机获取操作的数据结构。

先看三个操作的具体要求：

• 插入：insert(val)，向集合插入 val；
• 删除：remove(val)，val 存在时，移除 val；
• 随机获取：getRandom()，在集合中随机获取一个元素，其中返回的概率与元素在集合中的数量呈线性关系。

这里需要注意的是，题目允许出现重复的元素。

下面说一下三个操作的具体做法：

插入操作

这里我们直接使用数组的话，能够实现时间复杂度为 O(1)。

随机获取操作

这里只要随机获取数组中的一个索引就能够得到一个随机元素。

删除操作

数组的删除操作并不是 O(1)。题目中并没有说要保证元素的顺序，那么我们在删除的时候，先考虑将要被删除的元素元素与数组中的最后一个元素进行交换，那么就可以实现 O(1) 的时间复杂度。

那么这里，需要额外维护一个集合，存储数值在数组中每次出现的索引。

具体的代码实现如下。

class RandomizedCollection:

    def __init__(self):
        """
        Initialize your data structure here.
        """
        # 字典，key 为值，value 为值对应的索引
        self.ele_indexs = collections.defaultdict(set)
        self.elements = []


    def insert(self, val: int) -> bool:
        """
        Inserts a value to the collection. Returns true if the collection did not already contain the specified element.
        """
        # 添加索引到该值对应的索引中
        self.ele_indexs[val].add(len(self.elements))
        # 添加元素
        self.elements.append(val)
        # 示例中，如果元素已经在集合中，需要返回 False，这里判断下
        return len(self.ele_indexs[val]) == 1
        


    def remove(self, val: int) -> bool:
        """
        Removes a value from the collection. Returns true if the collection contained the specified element.
        """
        # 删除元素，先判断元素是否存在
        if not self.ele_indexs[val]:
            return False
        # 否则，待删除元素与集合末尾元素进行交换
        # 元素可能重复，取元素其中一个对应的索引
        i = self.ele_indexs[val].pop()
        # 这里直接用末尾元素替换
        self.elements[i] = self.elements[-1]
        # 这里要注意，末尾元素集合对应的索引已经改变，要注意移除，添加新的索引
        last_num = self.elements.pop()
        # 弹出 last_num 后，这里 self.elements 元素数量减 1
        self.ele_indexs[last_num].discard(len(self.elements))
        #
        if i < len(self.elements):
            self.ele_indexs[last_num].add(i)
        return True
        

    def getRandom(self) -> int:
        """
        Get a random element from the collection.
        """
        return random.choice(self.elements)


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

复杂度分析

• 时间复杂度：$O(1)$。
• 空间复杂度：$O(N)$，$N$ 表示集合元素数目。

欢迎关注

---

公众号 【书所集录】

---

如有错误，烦请指出，欢迎指点交流。