Problem 75 · Object & System Design

Randomized Set: O(1) Insert, Remove, GetRandom

MediumClassic DSA & systems

01Problem

Implement a set with insert(val), remove(val), and get_random() all in average O(1), where get_random returns each present element with equal probability.

02Key insight

Array for uniform random indexing plus a value->index map. To remove in O(1), swap the target with the last array element, update that element's index, and pop.

03Solution

solution.pypython30 lines

import random


class RandomizedSet:
    def __init__(self, rng: random.Random | None = None) -> None:
        self.items: list[int] = []
        self.pos: dict[int, int] = {}
        self.rng = rng or random.Random()

    def insert(self, val: int) -> bool:
        if val in self.pos:
            return False
        self.pos[val] = len(self.items)
        self.items.append(val)
        return True

    def remove(self, val: int) -> bool:
        i = self.pos.pop(val, None)
        if i is None:
            return False
        last = self.items.pop()
        if i < len(self.items):  # val wasn't the last element
            self.items[i] = last
            self.pos[last] = i
        return True

    def get_random(self) -> int:
        if not self.items:
            raise IndexError("get_random from empty set")
        return self.items[self.rng.randrange(len(self.items))]

Try it first.

TimeO(1) average per op

SpaceO(n)

tests.py · 4 passing

tests.pypython9 lines

rs = RandomizedSet(random.Random(0))
assert rs.insert(1) and not rs.insert(1) and rs.insert(2) and rs.insert(3)
assert rs.remove(1) and not rs.remove(1)
assert {rs.get_random() for _ in range(200)} == {2, 3}
assert rs.remove(3) and rs.remove(2)
try:
    rs.get_random(); assert False
except IndexError:
    pass

04Staff-level follow-up

Allow duplicates (LC 381)? Map value -> set of indices; the swap-with-last trick still works, updating one index in the moved value's set.