一.引言
设计一个支持在平均 时间复杂度 O(1) 下, 执行以下操作的数据结构。
注意: 允许出现重复元素。
集合包含以下三个功能:
A.insert(val):向集合中插入元素 val。
B.remove(val):当 val 存在时,从集合中移除一个 val。
C.getRandom:从现有集合中随机获取一个元素。每个元素被返回的概率应该与其在集合中的数量呈线性相关。
集合功能示例:
二.实现
集合类型定义为 RanomizedSet,实现了上述三个方法即插入元素,移除元素以及随机获取与元素个数线性相关的元素且时间复杂度为 O(1)。这里保存集合元素采用了 list,判断元素是否存在则使用了 dict 辅助,所以虽然该数据结构上述操作的时间复杂度为 O(1),但是为此空间复杂度也随之增加。
#!/usr/bin/python
# -*- coding: UTF-8 -*-
from random import choice
class RandomizedSet:
def __init__(self):
"""
Initialize your data structure here.
"""
self.dict = {}
self.list = []
def insert(self, val: int) -> bool:
"""
Inserts a value to the set. Returns true if the set did not already contain the specified element.
"""
self.list.append(val)
if val in self.dict:
return False
self.dict[val] = len(self.list) - 1
return True
def remove(self, val: int) -> bool:
"""
Removes a value from the set. Returns true if the set contained the specified element.
"""
if val in self.dict:
last_element, idx = self.list[-1], self.dict[val]
self.list[idx], self.dict[last_element] = last_element, idx
self.list.pop()
del self.dict[val]
return True
return False
def getRandom(self) -> int:
"""
Get a random element from the set.
"""
return choice(self.list
三.详解
1.insert(val)
def insert(self, val: int) -> bool:
"""
Inserts a value to the set. Returns true if the set did not already contain the specified element.
"""
self.list.append(val)
if val in self.dict:
return False
self.dict[val] = len(self.list) - 1
return True
向集合中插入元素,由于要求中提到允许出现重复的元素,所以这里开头就执行 list.append 的操作添加元素,这里原题目中给出的代码有问题,无法添加重复元素。List 中添加新元素 val 后,存储 val 在列表中对应的位置,这里通过 dict 保证一个元素只保存一个位置。返回值通过 dict 判断,如果 dict 存在则返回 False,否则返回 True。
2.remove(val)
def remove(self, val: int) -> bool:
"""
Removes a value from the set. Returns true if the set contained the specified element.
"""
if val in self.dict:
last_element, idx = self.list[-1], self.dict[val]
self.list[idx], self.dict[last_element] = last_element, idx
self.list.pop()
del self.dict[val]
return True
return False
移除元素中的元素,这里是这个 collection 最复杂的地方(其实也还好),这里的思路类似于冒泡排序一样。
A.获取最后一个元素的值和要删除值的索引
last_element, idx = self.list[-1], self.dict[val]
B.删除值位置的元素替换为列表最后一个值,并修改其在 dict 的索引位置
self.list[idx], self.dict[last_element] = last_element, idx
C.删除该值
上一个步骤修改列表 idx 处的值其实已经将 val 元素删除,这一步主要是更新数据结构,pop 减小列表长度并删除 dict 中多余的索引。
D.返回值
包含元素并移除返回 True,否则返回 False。
3.getRandom(val)
这里 choice 函数来自 random 库,需要 from random import choice 引入,这里如果想自己实现也很简单,只需要随机 [0, len(list)-1] 的数字并从 list 取值即可,随机到每个索引的概率为 1/n,n 为列表长度,所以 p(val) = 1/n * num,num 为该数字在集合中的个数,因此 p(val) 满足概率与其在集合中的数量呈线性相关。
def getRandom(self) -> int:
"""
Get a random element from the set.
"""
return choice(self.list)