The machine must hold a set with three operations, all average O(1): insert, remove, and getRandom — the last returning a uniformly random element. A plain set fails getRandom (no indexed access in O(1)); a plain array fails remove (O(n) deletion). The standard trick combines both: an array holds the elements, a dict maps each value to its array index.
insert appends the value to the array and records the index in the dict — O(1). remove is the interesting part: naive deletion would shift everything after the removed element, O(n). Instead, the machine swaps the victim with the last element, updates the dict entry for the moved element, pops the tail, and deletes the victim's key. Because order does not matter — this is a set, not a list — swap-and-pop preserves all invariants. getRandom picks a random index in [0, len) and returns the array value there.
The invariant that keeps everything correct: every array index 0..len−1 is occupied, and every dict value points at a live array position. A single misplaced dict entry after the swap silently corrupts later removes. The dict double-lookup is also a trap — after the swap the machine must update the dict of the element it just moved, which is the element that used to be at the tail.
Edge cases: removing the last element — the swap is a no-op, but the machine must still pop and delete the key; removing an absent value returns false; insert of a duplicate returns false and changes nothing; getRandom on an empty set is a caller error (or raise). All operations average O(1) with O(n) space.
import random
class RandomizedSet:
def __init__(self):
self.values = []
self.pos = {}
def insert(self, val):
if val in self.pos:
return False
self.pos[val] = len(self.values)
self.values.append(val)
return True
def remove(self, val):
if val not in self.pos:
return False
i = self.pos[val]
last = self.values[-1]
self.values[i] = last
self.pos[last] = i
self.values.pop()
del self.pos[val]
return True
def getRandom(self):
return random.choice(self.values)