Add, remove and pick at random: A Sample Coding Interview Solution

The coding question

Build a lottery system with add(person), remove(person) and pickRandom(). Every operation should run in constant time on average.

Ask a few questions first

Ask	What we'll assume
Can the same person be added twice?	No. Adding someone who's already in does nothing and returns False.
What does remove do for someone who isn't there?	Nothing, and returns False.
Does every person need an equal chance?	Yes, pickRandom is uniform over everyone currently in.
Does pickRandom remove the winner?	No, it just returns them. They stay in.
What if the lottery is empty?	Raise an error.

A couple of examples

Input	Output
add('ana'), add('ben'), add('cy'), add('ben')	True, True, True, False
then remove('ben'), remove('ben')	True, False, and the list is now ['ana', 'cy']
then pick_random()	'ana' or 'cy', each half the time

Start with brute force

Keep everyone in a list. Adding is an append, picking is a random index. But removing means finding the person first and then shifting everyone after them left.

Cost: remove is O(n). If you use only a set instead, remove is O(1) but picking a random item means walking the set, which is O(n).

A better way

Use both. A list gives you O(1) random picks by index. A hash map from person to their index in the list tells you where anyone is in O(1).

Removing from the middle of a list is the slow part, but removing from the end is fast. So move the last person into the removed person's spot, update their index in the map, and pop the end. Order doesn't matter in a lottery, so nothing's lost.

Keep a list of people and a map from person to list index.
add: if they're in the map, return False. Otherwise append to the list and store their index.
remove: look up their index, pop the last person off the list, and if the removed spot isn't the end, put the last person there and update their index.
pick_random: pick a random entry from the list.

The code

solution.pypython

import random


class Lottery:
    def __init__(self):
        self.people = []   # everyone in the draw, in no particular order
        self.index = {}    # person -> their position in self.people

    def add(self, person):
        if person in self.index:
            return False
        self.index[person] = len(self.people)
        self.people.append(person)
        return True

    def remove(self, person):
        if person not in self.index:
            return False
        spot = self.index.pop(person)
        last = self.people.pop()
        if spot < len(self.people):
            # Move the last person into the gap we just made.
            self.people[spot] = last
            self.index[last] = spot
        return True

    def pick_random(self):
        if not self.people:
            raise ValueError("nobody is in the lottery")
        return random.choice(self.people)

	Cost	Why
add	O(1) average	A map lookup and a list append
remove	O(1) average	A map lookup, a pop from the end and one swap
pick_random	O(1)	A random index into the list
Space	O(n)	Each person is in the list once and the map once

Test it

Walk through one of the examples out loud, then check the cases that tend to break things:

Removing the last person in the list: there's nobody to swap in, so just pop. The code checks spot < len(self.people) for this.
Removing the only person: the list and map both end up empty.
Adding someone after removing them: they get a fresh index at the end.
pick_random on an empty lottery: raises instead of crashing on an index error.

Common mistakes

Updating the moved person's index before popping, or forgetting to update it at all. The map then points at the wrong spot and a later remove deletes the wrong person.
Calling list.remove(person), which looks fine but is O(n).
Picking a random key from the map by turning it into a list first. That's O(n) every pick.
Not handling duplicates, so one person is in the list twice and has double the odds.

Follow-up questions

What if the same person can buy several tickets?

Store one list entry per ticket, and map each person to a set of their indexes. Removing one ticket swaps just that entry.

What if tickets have different weights?

Now picks aren't uniform. Keep running totals of weights and pick with a binary search on a random number, which is O(log n). Updates get harder, so a tree of partial sums helps.

How would you pick k distinct winners?

Pick one, remove them, repeat k times, then add them back if they should stay in. Each round is still O(1).

How the time adds up

Minutes	What you're doing
0–5	Clarify duplicates, empty draws and what pick returns
5–10	A list alone and a set alone, and what each gets wrong
10–20	Combine them and explain the swap-with-last trick
20–35	Write the class
35–40	Walk through adds and removes, including removing the last item
40–45	Follow-up questions

Key concepts used here

We wrote this sample solution ourselves. Practice Loop is independent and isn't affiliated with or endorsed by Bloomberg.