Short average O(1) Python beats 97%


  • 0

    The key idea is to use dict and set to achieve average O(1). And for the most difficult part in this problem, the remove, we swap the last elements with the val we wanna remove and update corresponding indices. Then just pop it and we are done.

    One thing to keep in mind is when updating self.indices[last], we should add before pop, otherwise we will have some trouble when i == j.

    class RandomizedCollection(object):
    
        def __init__(self):
            self.vals, self.indices = [], collections.defaultdict(set)
            
    
        def insert(self, val):
            self.indices[val].add(len(self.vals))
            self.vals.append(val)
            return len(self.indices[val]) == 1
            
    
        def remove(self, val):
            if self.indices[val]:
                last = self.vals[-1]
                i, j = len(self.vals) - 1, self.indices[val].pop()
                self.vals[i], self.vals[j] = self.vals[j], self.vals[i]
                self.vals.pop()
                self.indices[last].add(j)
                self.indices[last].remove(i)
                return True
            return False
    
    
        def getRandom(self):
            return random.choice(self.vals)
    

Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.