LC 0381 [H] Insert Delete GetRandom O(1) Duplicates allowed - ALawliet/algorithms GitHub Wiki

instead of a map of val to pos, we use a map of val to set(positions)

class RandomizedCollection:

    def __init__(self):
        self.hashmap = defaultdict(set) # val to set of positions instead of val to single pos
        self.list = []
        

    def insert(self, val: int) -> bool:
        self.hashmap[val].add(len(self.list))
        self.list.append(val)
        return len(self.hashmap[val]) == 1
        

    def remove(self, val: int) -> bool:
        if self.hashmap[val]:
            lastElement = self.list[-1]
            index = self.hashmap[val].pop()

            self.list[index] = lastElement
            self.hashmap[lastElement].add(index)

            self.hashmap[lastElement].discard(len(self.list) - 1)
            self.list.pop()

            return True
        return False
        

    def getRandom(self) -> int:
        return random.choice(self.list)
        


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