mirror of
https://github.com/heqin-zhu/algorithm.git
synced 2024-03-22 13:30:46 +08:00
63 lines
1.9 KiB
Python
63 lines
1.9 KiB
Python
#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 一次删除一个元素
|
|
#######################################################################
|
|
'''
|
|
|
|
|
|
class RandomizedCollection:
|
|
def __init__(self):
|
|
|
|
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
|
|
|
|
def removeAll(self, val: int) -> bool:
|
|
if val not in self.index:
|
|
return False
|
|
begin = end = len(self.vals)-len(self.index[val])
|
|
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
|
|
|
|
def remove(self, val):
|
|
if val not in self.index:
|
|
return False
|
|
last = len(self.vals)-1
|
|
idx = self.index[val].pop()
|
|
if len(self.index[val]) == 0:
|
|
del self.index[val]
|
|
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
|
|
|
|
def getRandom(self) -> int:
|
|
if self.vals:
|
|
return self.vals[random.randint(0, len(self.vals)-1)]
|