algorithm-in-python/dataStructure/insert_remove_getRandom.py

63 lines
1.9 KiB
Python
Raw Normal View History

#coding: utf-8
''' mbinary
#######################################################################
# File : insert_remove_getRandom.py
# Author: mbinary
# Mail: zhuheqin1@gmail.com
# Blog: https://mbinary.xyz
# Github: https://github.com/mbinary
# Created Time: 2019-05-24 10:01
# Description:
insert, remove, getRandom 的摊还时间为 O(1),
且有重复数据, remove 一次删除一个元素
#######################################################################
'''
2020-04-15 12:28:20 +08:00
class RandomizedCollection:
def __init__(self):
2020-04-15 12:28:20 +08:00
self.vals = []
self.index = {}
def insert(self, val: int) -> bool:
self.vals.append(val)
if val in self.index:
self.index[val].add(len(self.vals)-1)
return False
else:
self.index[val] = {len(self.vals)-1}
return True
2020-04-15 12:28:20 +08:00
def removeAll(self, val: int) -> bool:
if val not in self.index:
return False
begin = end = len(self.vals)-len(self.index[val])
2020-04-15 12:28:20 +08:00
for idx in self.index.pop(val):
if idx < begin:
while self.vals[end] == val:
end += 1
self.vals[idx] = self.vals[end]
self.index[self.vals[idx]].remove(end)
self.index[self.vals[idx]].add(idx)
self.vals = self.vals[:begin]
return True
2020-04-15 12:28:20 +08:00
def remove(self, val):
if val not in self.index:
return False
last = len(self.vals)-1
idx = self.index[val].pop()
2020-04-15 12:28:20 +08:00
if len(self.index[val]) == 0:
del self.index[val]
2020-04-15 12:28:20 +08:00
if idx != last:
self.vals[idx] = self.vals[last]
self.index[self.vals[idx]].remove(last)
self.index[self.vals[idx]].add(idx)
self.vals.pop()
return True
2020-04-15 12:28:20 +08:00
def getRandom(self) -> int:
if self.vals:
2020-04-15 12:28:20 +08:00
return self.vals[random.randint(0, len(self.vals)-1)]