algorithm-in-python/dataStructure/insert_remove_getRandom.py

#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)]