Problem statement

https://leetcode.com/problems/insert-delete-getrandom-o1-duplicates-allowed/

Solution

The idea is to keep defaultdict with set of indexes and list with all our numbers.

Example: lst = [1,2,1,3,3,1,4], idx = {1:[0,2,5], 2:[1], 3:[3,4], 4:[6]}.

For getRandom we just use random choice of elements of list in O(1). For insert(val) we add new index of element to our defaultdict and add it to the list, return True if element is new, that is size of set of indexes for this element is equal to 1. The trickiest part is remove. Let use define indexes remove, last = self.idx[val].pop(), self.lst[-1]. Then we need to update our lst: instead of gap put our removed element: self.lst[remove] = last and remove the last element. Also we need to update our self.idx[last]: we add remove and remove len(self.lst). Try to do it on given example, and remove element 1.

Complexity

It is O(1) both for

Code

class RandomizedCollection:
    def __init__(self):
        self.lst = []
        self.idx = defaultdict(set)

    def insert(self, val):
        self.idx[val].add(len(self.lst))
        self.lst.append(val)
        return len(self.idx[val]) == 1

    def remove(self, val):
        if not self.idx[val]: return False
        
        remove, last = self.idx[val].pop(), self.lst[-1]
        self.lst[remove] = last
        self.lst.pop()
        
        self.idx[last].add(remove)
        self.idx[last].remove(len(self.lst))
        return True

    def getRandom(self):
        return choice(self.lst)