Insert, delete and random pick in O(1) in Python

Why this matters

A set is a collection that holds each value at most once. You can ask it three things: add a value, remove a value, and check whether a value is inside. A Python set answers all three in about the same time no matter how many values it holds. The time for each call stays roughly the same as the set grows. Programmers call this O(1), which reads as "constant time".

A set cannot do one more thing: give you a random value from it. It has no position numbers. You cannot ask it for "item number 3", so picking a random item means walking through the whole set first.

A random pick is useful in real programs. A game might choose a random enemy. A test might pick a random sample from a list of test cases. A simple load balancer might choose a random server. Each of these needs a collection where you can add, remove and pick at random, all fast.

This lesson builds that structure in Python. It uses two standard types side by side, a list and a dict, and one trick that keeps removal fast. You will see the exact output of each program, including a mistake that breaks the structure and the error it produces.

The problem in plain words

Design a class called RandomizedSet with three methods:

insert(value) adds a value. It returns True if the value was added and False if the set already had it.
remove(value) removes a value. It returns True if the value was there and False if not.
get_random() returns one value from the set. Every value in the set must have the same chance of being returned.

All three methods must take constant time on average. Here, "on average" means that the time per call stays about the same over many calls, even if a few calls take longer. Python's dict works this way too.

Here is the behaviour in one example. The set starts with the values 10, 20 and 30. Then remove(10) runs. Afterwards the set holds 20 and 30 only. A call to get_random() returns 20 half of the time and 30 half of the time.

Before building it, look at why two simpler ideas fail.

A set gives fast add, remove and in, but it cannot pick a random item quickly.

A list gives fast random pick, because values[index] jumps straight to any position. Adding at the end is also fast. The problem is removal. If you remove the item at position 0 from [10, 20, 30], Python shifts the later items one place to the left, so the list becomes [20, 30]. With a long list, every removal near the front moves almost every item. Also, in on a plain list searches from the start, which takes time that grows with the list size.

The structure in this lesson combines the two. The list gives random access. A dict gives fast lookup. The swap trick stops removal from shifting items.

The idea in two parts

The structure has two parts that always agree with each other.

A list called values holds the items in order. Position 0 is the first item, position 1 is the second item, and so on. Positions are called indexes. The first index is 0.
A dict called positions holds each value as a key, and its current index in the list as the value. Looking up a key in a dict is fast.

When a value is stored, it is in the list at one index, and the dict says that index. Keep that promise, and the rest follows.

Insert

To insert a value, first check the dict. If the value is already a key in the dict, the set already has it, so return False. Otherwise append the value to the end of the list, and record its new index in the dict.

Here is the insert on the set [10, 20]. The dict is {10: 0, 20: 1}. Insert 30:

30 is not a key in the dict, so it is new.
Append 30 to the list. The list becomes [10, 20, 30]. Its index is 2.
Store 30 with the value 2 in the dict. The dict becomes {10: 0, 20: 1, 30: 2}.

Appending to the end of a list is fast, and so is a dict store. So insert is constant time on average.

Random pick

To pick a random value, choose a random index from 0 up to the list length minus one, then return the value at that index. The random.randrange(n) function returns a number from 0 up to n - 1, with each number equally likely. Since every index is equally likely, every value is equally likely.

On the list [30, 20], randrange(2) returns 0 or 1. Index 0 gives 30 and index 1 gives 20. Each has probability one half.

Picking takes constant time, because values[index] jumps straight to the position.

Remove, the hard part

Removing is where the trick is needed. Here is the set [10, 20, 30] with the dict {10: 0, 20: 1, 30: 2}. Remove 10.

Step 1. Look up 10 in the dict. It is at index 0.

Step 2. Look at the last item in the list. The last index is 2, and the value there is 30.

Step 3. Copy the last value into the slot that is being removed. Position 0 now holds 30, so the list is [30, 20, 30]. Right now 30 appears twice, but that is only temporary.

Step 4. Remove the last slot. The list is [30, 20].

Step 5. Update the dict. The value 30 moved from index 2 to index 0, so the dict must say 30: 0. The value 10 is gone, so remove key 10. The dict is now {20: 1, 30: 0}.

The final state is correct. The list holds 30 and 20 at indexes 0 and 1. The dict agrees.

Why swap instead of shifting? Shifting moves every item after the removed one, so the time grows with the list size. The swap changes one slot and then removes the last slot. Removing the last item of a Python list moves nothing. So removal is constant time.

The swap changes the order of the list. That does not matter here, because a set has no order. You only care about which values are inside.

The last item as a special case

Look at what happens when the value being removed is already the last item. Take the list [10, 20, 30] and remove 30.

The dict says 30 is at index 2.
The last index is also 2, and its value is 30.
Copy 30 into index 2. Nothing changes.
Update the dict: 30: 2. Nothing changes.
Remove the last slot. The list is [10, 20].
Remove key 30 from the dict.

The same five steps work with no special case. Keep the order of the steps. The dict update for the moved value must come before the dict removal of the removed value. If the value is the last item, both point at the same key, and the removal must come last so the key is gone at the end.

The first program

The program below builds the set and runs a short script. Each line of output comes from one call. The __repr__ method prints the list, so you can see the state after each change.

Python · editable

import random

class RandomizedSet:
    def __init__(self, seed=None):
        self.values = []
        self.positions = {}
        self.rng = random.Random(seed)

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

def remove(self, value):
        if value not in self.positions:
            return False
        index = self.positions[value]
        last_index = len(self.values) - 1
        last_value = self.values[last_index]
        self.values[index] = last_value
        self.positions[last_value] = index
        self.values.pop()
        del self.positions[value]
        return True

def contains(self, value):
        return value in self.positions

def get_random(self):
        if not self.values:
            raise ValueError("set is empty")
        return self.values[self.rng.randrange(len(self.values))]

def __repr__(self):
        return repr(self.values)

s = RandomizedSet(42)
print("insert 10:", s.insert(10))
print("insert 20:", s.insert(20))
print("insert 30:", s.insert(30))
print("insert 20 again:", s.insert(20))
print("state:", s)
print("remove 10:", s.remove(10))
print("state:", s)
print("remove 99:", s.remove(99))
print("contains 10:", s.contains(10))
print("contains 30:", s.contains(30))
print("remove 20 (last item):", s.remove(20))
print("state:", s)

picks = RandomizedSet(7)
for v in [1, 2, 3, 4]:
    picks.insert(v)
print("five random picks:", *[picks.get_random() for _ in range(5)])

try:
    RandomizedSet(1).get_random()
except ValueError as error:
    print("empty set:", error)

insert 10: True
insert 20: True
insert 30: True
insert 20 again: False
state: [10, 20, 30]
remove 10: True
state: [30, 20]
remove 99: False
contains 10: False
contains 30: True
remove 20 (last item): True
state: [30]
five random picks: 3 2 4 1 1
empty set: set is empty

Read the output line by line. The state line after remove 10 shows [30, 20], which is the swap from the previous section. The random picks use the seed 7, so the same five numbers come out every time you run the program on the same Python version. A seed is the starting number of a random number generator. The same seed gives the same sequence. The class keeps its own random.Random object, so one set's seed does not change another set's picks.

Notice that the picks from [1, 2, 3, 4] include 1 twice. That is normal for random picks. Each pick is independent, so a value can repeat. With five picks from four values, a repeat is certain.

The second program: checking against a trusted set

How do you know the structure is correct? One method is to run random operations on your structure and on a trusted set at the same time, then compare the answers. This is called a stress test. It finds mistakes that a short script would miss, because the random sequence reaches many states.

The program below does 200,000 random operations. Each one is an insert, a remove or a contains check, on values from 0 to 49. After each operation it compares the size of both collections. Any difference counts as a mismatch.

Python · editable

import random

class RandomizedSet:
    def __init__(self, seed=None):
        self.values = []
        self.positions = {}
        self.rng = random.Random(seed)

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

def contains(self, value):
        return value in self.positions

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

rs = RandomizedSet(2026)
reference = set()
operations = 200000
mismatches = 0

for step in range(operations):
    value = rs.rng.randrange(50)
    action = rs.rng.randrange(3)
    if action == 0:
        if rs.insert(value) != (value not in reference):
            mismatches += 1
        reference.add(value)
    elif action == 1:
        if rs.remove(value) != (value in reference):
            mismatches += 1
        reference.discard(value)
    else:
        if rs.contains(value) != (value in reference):
            mismatches += 1
    if len(rs.values) != len(reference):
        mismatches += 1

picks_in_set = 0
for _ in range(1000):
    if rs.get_random() in reference:
        picks_in_set += 1

print("operations:", operations)
print(f"final size: {len(rs.values)} (reference {len(reference)})")
print("mismatches:", mismatches)
print("random picks that were in the set:", picks_in_set, "of 1000")

operations: 200000
final size: 24 (reference 24)
mismatches: 0
random picks that were in the set: 1000 of 1000

Zero mismatches means every insert, remove and contains answer matched the trusted set, and the size matched after every step. The last line checks that get_random only returns values that are really in the set.

The seed is fixed, so the numbers are the same each time you run this program on the same Python version. A different version of Python may give different numbers, but the mismatch count must still be zero.

Third example: what goes wrong without the dict update

The most common bug in this structure is forgetting to update the dict for the value that moved. The code below does the swap but skips the dict update, then reads back the dict. This is a deliberate mistake. Read the output, then find the step that is missing.

values = [10, 20, 30]
positions = {10: 0, 20: 1, 30: 2}
 
# remove 10 without fixing the moved item
index = positions[10]
values[index] = values[-1]
values.pop()
del positions[10]
 
print("values:", values)
print("positions:", positions)
print(values[positions[30]])

values: [30, 20]
positions: {20: 1, 30: 2}
Traceback (most recent call last):
  ...
IndexError: list index out of range

The list says 30 is at index 0. The dict still says 30: 2, because nothing updated it. Looking up index 2 in a list of length 2 fails. The program stopped at the last print.

The fix is the line positions[last_value] = index in the remove method of the first program. Without it, every later call that uses the dict can point at the wrong slot or at a slot that no longer exists.

Common mistakes

Forgetting to update the moved value in the dict. The symptom is the exception shown in the third example. A less obvious symptom is remove raising IndexError: list assignment index out of range for a value that is still in the set. Every time you move a value in the list, write its new index to the dict.

Deleting the removed key too early. Take the case where the removed value is the last item. If you delete its dict entry first and then write the moved value's new index, the entry comes back with an index that no longer exists. Keep the order: move the last value, update its dict entry, shrink the list, then delete the removed key.

Using remove on a list when you meant pop. A Python list has two removal methods, and they take different things. list.remove(value) looks for the value and removes its first match. list.pop(index) removes the item at a position. The structure in this lesson needs pop() with no argument, which removes the last item. Mixing them up gives wrong results or an error:

Python · editable

after remove(10): [20]

y = [10, 20]
y.pop(10)

IndexError: pop index out of range

The first call removed the value 10, and the second tried to remove the item at position 10, which does not exist. Read the method name and its argument before you rely on the result.

Calling get_random on an empty set. There is no value to return. The class above raises ValueError with the message set is empty. A caller should check len(values) first, or handle the exception.

Using a set and expecting a random pick. A set has no index, so there is no position to choose. The structure in this lesson exists because of this limit. random.choice(some_set) fails with a TypeError, because a set does not support indexing, and random.choice needs one.

Comparing with is for values. Small integers are shared objects in CPython, so is can look right for them and fail for larger values. Use == to compare values, as the stress test does.

Practice

Exercise 1. The list is [5, 6, 7, 8] and the dict is {5: 0, 6: 1, 7: 2, 8: 3}. Call remove(6). What are the list and the dict afterwards?

Show the answer

The list is [5, 8, 7]. The dict is {5: 0, 8: 1, 7: 2}. The value 8 was the last item, so it moved to index 1. Its dict entry changed from 8: 3 to 8: 1. The key 6 was removed from the dict.

Exercise 2. The set holds 1, 2 and 3 in that order, so the list is [1, 2, 3]. Run these calls in order: remove(1), contains(3), remove(2), then get_random(). What does each call return, and what is the list at the end?

Show the answer

remove(1) returns True. The list becomes [3, 2]. contains(3) returns True. remove(2) returns True. The list becomes [3]. get_random() returns 3, because it is the only value left. The final list is [3].

Exercise 3. Why is the dict update positions[last_value] = index needed, even though last_value is still in the list at the end?

Show the answer

The value moved to a new index, so its dict entry must change. Without the update, the dict still points at the old index, which is the last slot. That slot is gone after the list shrinks, so a later lookup fails. The third example in this lesson shows the error that this bug causes.

Exercise 4. Explain in your own words why swapping with the last item is constant time, but shifting every later item left is not.

Show the answer

Swapping changes one slot in the list and then removes the last slot. Removing the last slot moves no other item. Shifting moves every item after the removed one, one place to the left. A list with a million items would move almost a million items to remove one value near the front. The shifting time grows with the size of the list, so it is not constant.

Exercise 5. Change get_random so it returns None for an empty set, instead of raising an error. What would the caller need to check?

Show the answer

Add if not self.values: return None at the top of the method. The caller must check for None before using the result, because using None as a number or as an object raises TypeError or AttributeError later, far from the real cause. Raising an error makes the mistake show up at the call that caused it, which is one reason many programmers prefer the error.

Summary

A set finds and removes values quickly, but it cannot pick a random value, because it has no positions.
The structure keeps a list of values for random picks, and a dict from each value to its list index for fast lookup.
Removal swaps the removed value with the last item, removes the last slot, and updates the dict for the moved value.
The swap avoids shifting many items, so removal stays constant time on average.
A stress test that compares your structure with a trusted set finds bugs that short examples miss.

Insert, delete and random pick in O(1) in Python

Why this matters

The problem in plain words

The idea in two parts

Insert

Random pick

Remove, the hard part

The last item as a special case

The first program

The second program: checking against a trusted set

Third example: what goes wrong without the dict update

Common mistakes

Practice

Summary

Stuck? Ask the tutor