DSA Insert, Delete and GetRandom in O(1)

The problem

The spec asks for a set with three operations. Insert adds a value and reports whether it was new. Remove deletes a value and reports whether it was present.

The third operation returns a random member. Every stored value must have the same chance of being drawn. All three operations must run in constant time, which means their step count does not grow with the size of the set.

script = [
    ("insert", 10, True),
    ("insert", 20, True),
    ("insert", 30, True),
    ("insert", 40, True),
    ("insert", 10, False),
    ("remove", 20, True),
    ("remove", 99, False),
    ("get_random", None, "one of the stored values"),
]
for op, value, expected in script:
    print(op, value, "->", expected)

The script fixes the test for the whole article. The naive design and the real design both run it. Their step counts are compared at the end.

Intuition: one big list on its own

The first design is a single list. A random member is easy, because a random index picks one. Insert and remove carry the hidden costs.

Insert must first check that the value is absent, and that check scans the list. Remove must close the gap it leaves, and closing a gap shifts every later value one slot left.

values = [10, 20, 30, 40]
checks = 0
for value in values:
    checks += 1
    if value == 30:
        break
print("finding 30 for a duplicate check took", checks, "comparisons")
hole = values.index(20)
shifts = len(values) - hole - 1
print("removing 20 leaves a hole at index", hole, "and shifts", shifts, "values")

Both costs grow with the list. A set of a thousand values makes the duplicate check scan up to a thousand entries. The design needs a way to jump straight to a value.

Brute force: the naive design, counted

The naive design is the single list, written as a class. A counter records its comparisons and shifts. The script from section 1 runs against a large set, where the costs are visible.

import random

class NaiveSet:
    def __init__(self, seed):
        self.values = []
        self.rng = random.Random(seed)
        self.work = 0

    def insert(self, value):
        for existing in self.values:
            self.work += 1
            if existing == value:
                return False
        self.values.append(value)
        return True

    def remove(self, value):
        for i, existing in enumerate(self.values):
            self.work += 1
            if existing == value:
                self.work += len(self.values) - i - 1
                self.values.pop(i)
                return True
        return False

    def get_random(self):
        self.work += 1
        return self.values[self.rng.randrange(len(self.values))]

bag = NaiveSet(7)
for value in range(1000):
    bag.insert(value)
bag.work = 0
print("insert a duplicate:", bag.insert(500), "work:", bag.work)
bag.work = 0
print("remove the first value:", bag.remove(0), "work:", bag.work)
bag.work = 0
print("one draw:", bag.get_random(), "work:", bag.work)

The draw was already constant time. The other two operations are not. Their work counts scale with the set, which is exactly what the spec forbids.

The optimal solution: a list for draws, a dict for positions

The state is two structures kept in step. The list holds the values for random draws, and a dict, a table that maps each value to a stored index, holds each value's position in the list.

Removal uses the position to jump to the hole. The last value moves into the hole, and its dict entry is updated. No value after the hole moves, so nothing shifts.

import random

class RandomSet:
    def __init__(self, seed):
        self.values = []
        self.where = {}
        self.rng = random.Random(seed)

    def insert(self, value):
        if value in self.where:
            return False
        self.where[value] = len(self.values)
        self.values.append(value)
        return True

    def remove(self, value):
        if value not in self.where:
            return False
        hole = self.where[value]
        last = self.values[-1]
        self.values[hole] = last
        self.where[last] = hole
        self.values.pop()
        del self.where[value]
        return True

    def get_random(self):
        return self.values[self.rng.randrange(len(self.values))]

bag = RandomSet(7)
print("insert 10:", bag.insert(10), bag.values)
print("insert 20:", bag.insert(20), bag.values)
print("insert 30:", bag.insert(30), bag.values)
print("insert 40:", bag.insert(40), bag.values)
print("insert 10 again:", bag.insert(10), bag.values)
print("remove 20:", bag.remove(20), bag.values)
print("positions now:", sorted(bag.where.items()))
print("remove 99:", bag.remove(99), bag.values)
print("seeded draws:", [bag.get_random() for _ in range(4)])

Every line of the script matches the expected results from section 1. The draws come from a generator seeded with 7 inside the code. The same seed draws the same sequence on every run.

Walkthrough: removing 20, step by step

The walk rebuilds the set and removes 20 slowly. The list and the positions are printed after each stage. Positions print as a sorted list of pairs, so the order never depends on the dict.

import random

class RandomSet:
    def __init__(self, seed):
        self.values = []
        self.where = {}
        self.rng = random.Random(seed)

    def insert(self, value):
        if value in self.where:
            return False
        self.where[value] = len(self.values)
        self.values.append(value)
        return True

    def remove(self, value):
        if value not in self.where:
            return False
        hole = self.where[value]
        last = self.values[-1]
        self.values[hole] = last
        self.where[last] = hole
        self.values.pop()
        del self.where[value]
        return True

    def get_random(self):
        return self.values[self.rng.randrange(len(self.values))]

bag = RandomSet(7)
for value in [10, 20, 30, 40]:
    bag.insert(value)
print("values:", bag.values)
print("positions:", sorted(bag.where.items()))
hole = bag.where[20]
last = bag.values[-1]
print("hole index:", hole, " last value:", last)
bag.values[hole] = last
bag.where[last] = hole
print("after moving the last value:", bag.values)
print("positions:", sorted(bag.where.items()))
bag.values.pop()
del bag.where[20]
print("after dropping the last slot:", bag.values)
print("positions:", sorted(bag.where.items()))

The value 40 moved from index 3 to index 1. Its dict entry moved with it. A later removal of 40 will jump straight to index 1.

Complexity

Each operation does a fixed number of dict lookups, list writes and pops, so each runs in O(1) time. The list and the dict together store one entry per value, so the space is O(n).

The proof is a step count at two sizes. The same operations run on a set of 5 values and a set of 1,000 values. The counts must not move.

import random

class CountedSet:
    def __init__(self, seed):
        self.values = []
        self.where = {}
        self.rng = random.Random(seed)
        self.steps = 0

    def insert(self, value):
        self.steps += 1
        if value in self.where:
            return False
        self.where[value] = len(self.values)
        self.values.append(value)
        self.steps += 1
        return True

    def remove(self, value):
        self.steps += 1
        if value not in self.where:
            return False
        hole = self.where[value]
        last = self.values[-1]
        self.values[hole] = last
        self.where[last] = hole
        self.values.pop()
        del self.where[value]
        self.steps += 3
        return True

    def get_random(self):
        self.steps += 1
        return self.values[self.rng.randrange(len(self.values))]

small = CountedSet(7)
for value in [10, 20, 30, 40]:
    small.insert(value)
small.steps = 0
print("insert 50 into 4 values:", small.insert(50), "steps:", small.steps)
small.steps = 0
print("remove 20 from 5 values:", small.remove(20), "steps:", small.steps)
small.steps = 0
print("one draw from 4 values:", small.get_random(), "steps:", small.steps)
big = CountedSet(7)
for value in range(1000):
    big.insert(value)
big.steps = 0
print("remove 500 from 1000 values:", big.remove(500), "steps:", big.steps)
big.steps = 0
print("one draw from 999 values:", big.get_random(), "steps:", big.steps)

Remove costs four steps at both sizes. The draw costs one step at both sizes. That flat count is the constant-time claim, shown rather than asserted.

Edge cases and pitfalls

Removing the last value is the delicate case. The hole is the last slot itself, so the value moves onto itself before the pop. The code handles it with no special branch.

A draw assumes the set is not empty. Calling it on an empty set has no member to return. The script in this article never draws from an empty set.

import random

class RandomSet:
    def __init__(self, seed):
        self.values = []
        self.where = {}
        self.rng = random.Random(seed)

    def insert(self, value):
        if value in self.where:
            return False
        self.where[value] = len(self.values)
        self.values.append(value)
        return True

    def remove(self, value):
        if value not in self.where:
            return False
        hole = self.where[value]
        last = self.values[-1]
        self.values[hole] = last
        self.where[last] = hole
        self.values.pop()
        del self.where[value]
        return True

    def get_random(self):
        return self.values[self.rng.randrange(len(self.values))]

bag = RandomSet(7)
for value in [10, 20, 30, 40]:
    bag.insert(value)
print("remove the last value 40:", bag.remove(40), bag.values)
print("positions:", sorted(bag.where.items()))
solo = RandomSet(7)
solo.insert(42)
print("remove the only value:", solo.remove(42), solo.values)
print("reinsert it:", solo.insert(42), solo.values)
print("draws from a one-value set:", [solo.get_random() for _ in range(3)])
print("duplicate insert:", solo.insert(42))
print("missing remove:", solo.remove(99))

The one-value draws print 42 three times. With one stored value, every random index is index 0. Randomness is only visible when the set holds a choice.