• Random sampling question, version 1
  • Monte Carlo solution (aka rejection sampling)
  • Naive deletion
  • Keeping track of gaps?
  • Swapping to the end
  • Via partial shuffle
    • A fun side quest
  • Floyd’s algorithm
• Version 2: \(O(N)\) solutions
  • Streams of known length
    • Selection sampling
  • Streams of unknown length
    • Heaps
    • Quickselect
    • Reservoir sampling
    • Reservoir sampling with skipping ahead
• Online shuffling
  • Block ciphers
  • Pseudorandom permutations
  • Is 0% enough?
  • Cheap parameterized permutations

I’m going to describe my favorite interview question and then ad-lib a bit on it. The question has a somewhat mathematical bent but connects to several practical problems I had in my career. I used it many dozens of times when interviewing for xAI and other places, to the extent that it’s slightly worn out (apparently it was asked at Google going back to at least 2005 and unsurprisingly is on actual leetcode). However, as per the last blog post, I don’t think I am giving “too much away” by describing it here, and I may well use this question again. If candidates read this, all the better – we can dive in deep then. And if they have not, despite knowing I’d interview them, that’s an ever so weak signal too.

I first heard this question when Greg Chanan interviewed me for FAIR in 2018. I probably did okay-but-not-great. I did end up getting an offer and I did join Facebook (renamed to Meta while I was there) in early 2019.

Without further ado, here’s the question in its first variant.

Random sampling question, version 1

Given an array L of length N and a number k with 0 <= k <= N, sample k elements from L, uniformly and without replacement.

I typically would suggest some kind of source of randomness, e.g. the Python standard library functions random.random and random.randrange, with the following behavior:¹

random.random()

Return the next random floating-point number in the range 0.0 <= X < 1.0

 

random.randrange(stop)

Return a randomly selected element from range(stop).

Some candidates want an explanation of the terms: “uniform” means every entry of the array has the same chance of being sampled; “without replacement” means elements are discarded after being sampled and cannot be sampled again: in the “urn model”, the balls taken out of the urn are not being put back into the urn, so they are not being replaced.

This opens the door to another question: What is a “same element” for the purpose of this question? It’s overall more interesting if there’s no object equality check a la __eq__ in Python here but instead all elements of the array can be assumed to be unique. This also implies their index is the only thing that matters, and what the elements of the array are is otherwise not important.

Monte Carlo solution (aka rejection sampling)

A first attempt to solve this could look like²

from random import randrange

def sample(L, k):
    result = set()
    while len(result) < k:
        result.add(L[randrange(len(L))])
    return result

This is decent if k/N is a small number but catastrophically slow when k is close to N. However, in the case k/N >= 0.5 one can choose to sample rejected instead of selected items and arrive at a decent solution overall. This way, the algorithm makes progress in each iteration with probability \(\ge 0.5\), which means the chance of not making progress after \(m\) steps decays exponentially in \(m\). Note though that the worst case time complexity is still unbounded (“\(O(\infty)\)”).

That’s good enough for most practical applications – but not good enough for a coding interview. So let’s say we are looking for a solution with deterministic complexity.

Naive deletion

One solution that maps well to “don’t replace balls in the urn” is simply

def sample(L, k):
    result = []
    for _ in range(k):
        j = randrange(len(L))
        result.append(L[j])
        del L[j]
    return result

In place of del L[j], some people even go for L = L[: j] + L[j + 1 :]. If they have more Python knowledge they might propose

def sample(L, k):
    result = []
    for _ in range(k):
        result.append(L.pop(randrange(len(L))))
    return result

All of these are correct solutions, but they suffer from a similar problem. In order to delete an element of an array while keeping \(O(1)\) lookups in place, deletion requires shifting all later elements which is \(O(N)\).

In memory, an array is a buffer that starts at start_ptr and ends at start_ptr + buffer_len:

[ a b c d e f g h i j k l m ]
  ^                         ^
  start_ptr                 |
       start_ptr + buffer_len

Looking up the element at index i means resolving a pointer to start_ptr + i. This makes lookups \(O(1)\) and cheap even within that class.

Deleting an element – say, g – creates a hole:

[ a b c d e f _ h i j k l m ]

The only way to keep simple \(O(1)\) lookups is to shift later elements into the hole:

[ a b c d e f h _ i j k l m ]
                  <

[ a b c d e f h i _ j k l m ]
                    <

[ a b c d e f h i j _ k l m ]
                      <

etc.

This makes deletion shift \(N/2\) elements on average and \(O(N)\).

Note that a Python list is also just such a buffer in memory with a physical capacity and a logical length. When it needs to grow beyond its capacity a new buffer is created and all elements (which are pointers of type PyObject*) are copied. Candidates often think they don’t know that that is the case but they typically can infer this once the “contract” of the list is pointed out, which is that random access is \(O(1)\). There is simply no alternative to an array under this contract! See listobject.c in the CPython source for the details.

Keeping track of gaps?

Some people, myself included, find it natural to try to track already selected elements. In step i, one could attempt to sample from randrange(N - i) and then add the number of gaps already chosen left of i. While natural, this is complicated and yields at best some kind of tree structure and an \(O(k \log k)\) solution, while \(O(k)\) should be and is possible.

Some folks propose all kinds of dictionary/hash map approaches, most of which fail (with an exception, not realistically findable in an interview, see below!). The fundamental problem is that arrays allow \(O(1)\) lookup of the \(i\)th element but require \(O(N)\) deletions, while hash maps might allow for \(O(1)\) deletions but looking up the \(i\)th element costs \(O(N)\). Hash maps generally are somewhat overused due to their convenience. Some candidates have a really hard time giving them up for this problem.

Swapping to the end

Still, modifying the array is promising, and candidates often find a better solution when they are explicitly told that’s an okay thing to do. For instance, one could do

from random import randrange

def sample(L, k):
    N = len(L)
    result = []
    for i in range(k):
        j = randrange(N - i)
        # Take element j, then move in element from the end.
        result.append(L[j])
        L[j] = L[N - i - 1]
    return result

Note that L[j] = L[N - i - 1] can be a no-op if the sampled index j was the end.

This is an example of a “swap with the end” trick that’s useful more generally. With swapping instead of moving one can also avoid the result array:

from random import randrange

def sample(L, k):
    N = len(L)
    for i in range(k):
        j = randrange(N - i)
        L[j], L[N - i - 1] = L[N - i - 1], L[j]
    return L[N - k :]

This is pretty short and \(O(k)\) in time with essentially no extra space, assuming the slice at the end doesn’t create a copy – true if L is an object like a NumPy array but not true for plain Python lists. If k is nonzero, the slice at the end can be written as L[-k :] but that fails for k == 0.

Via partial shuffle

Sometimes candidates find that “shuffle the list and take the first k items” is a more intuitive solution:

def sample(L, k):
    shuffle(L)
    return L[:k]

The question then is how to shuffle. A moment’s reflection shows that “go through all slots and swap with a random one” is promising:

def shuffle(L):
    for i in range(len(L)):
        j = i + randrange(len(L) - i)
        L[i], L[j] = L[j], L[i]

The logic of this algorithm (known as Fisher-Yates in the literature) is to keep a range of not-yet-chosen elements that is initially the full array and then shrinks from the left by one element each iteration. (Since the last random “choice” is choosing from the one-element range(1), the loop could actually use range(len(L) - 1).)

This has the nice property that each to-be-filled slot is chosen once and then the algorithm sticks to its choice.

It can be turned into a “partial shuffle” where one stops after k slots:

def shuffle(L, k):
    for i in range(k):
        j = i + randrange(len(L) - i)
        L[i], L[j] = L[j], L[i]

With this, shuffle(L, k) and then returning L[: k] is equivalent to the swapping-to-the-end solution, but easier to find for some people.

A fun side quest

Here’s a side question: What if we don’t stick with our choices, but instead allow the full array to be selected every time? This way we potentially swap already chosen elements:

def alt_shuffle(L):
    for i in range(len(L)):
        j = randrange(len(L))
        L[i], L[j] = L[j], L[i]

As an algorithm, this feels worse – for instance, the original shuffle allows us to yield results as they come. But is it any worse otherwise? Is it correct? Are both shuffle and alt_shuffle correct, or only one of them? Or, less plausibly, neither?

Floyd’s algorithm

What is optimal? In a sense, the partial shuffle/swapping-based solution is optimal – it’s \(O(k)\) with minimal space overhead. Its only imperfection is the fact that the input array L is modified. Of course this can be avoided by making a copy, but that means we are no longer \(O(k)\).

Can we avoid modifying the input? No simple change to the swapping-based solution will achieve that, because the algorithm uses the buffer as a scratchpad to keep track of which of the \(N!\) permutations has been selected.

An \(O(k)\) space and time solution to this exercise could dispense with the array altogether, since it really just selects \(k\) indices. One could simply ask:

Generate k distinct uniformly random integers in range(N) in \(O(k)\) space and time.

I had asked many candidates the sampling question before I found this version, innocently stated as Exercise 5-2 of the AWK programming book. Their proposed solution, translated to Python, reads like this:

from random import randrange

def sample(N, k):
    result = set()
    for i in range(N - k, N):  # k steps.
        r = randrange(i + 1)
        if r in result:
            r = i
        result.add(r)
    return result

This is known as Floyd’s algorithm, due to R.W. Floyd, 1978 Turing Award winner and close collaborator of Donald Knuth at Stanford.

It wasn’t immediately obvious to me why this algorithm is correct. The fact that the math of “if seen already, choose the end” works is somewhat magical. For a discussion of Floyd’s algorithm, check out “A Sample of Brilliance” in Jon Bentley’s Programming Pearls.

Very few people will be able to derive something like that in the scope of a coding interview. Only a handful of people even know this algorithm in the first place; nobody I ever interviewed did.

Version 2: \(O(N)\) solutions

One of the neat properties of this problem is its many variations. Having solved the original version, what if we look to sample from a stream instead? A stream is an object that we can ask for its next value, which will either yield another value or tell us the stream has ended. In Python, that maps to the concept of an iterator.

For this exercise to make sense, we need some guarantees: The stream needs to (1) have at least k elements, and (2) eventually end.

The point of asking about streams is to force the solution to be \(O(N)\), which gives a different landscape of trade-offs. Of course, we could simply collect all \(N\) items in an array and apply the solutions we already saw, but perhaps we can do this in \(O(k)\) space?

Streams of known length

Selection sampling

Let’s say we also get information about the eventual length of the stream. In that case, sampling \(k\) out of the \(N\) items we will see, naively, is³

from random import random

def sample(it, N, k):
    for e in it:
        if random() < k/N:
            yield e

The issue with this algorithm is that it could sample both fewer and more than \(k\) elements. It’s correct “on average” but each run has some variation.

Sampling too much is easily fixable – we can just stop after \(k\) samples. But it’s unclear if that maintains uniformity, and anyway what can we do about sampling too little? What we should do is to increase the probability as we go along, for instance in a situation where we have sampled \(k - 1\) elements and the stream only contains a single final element, that one has to be sampled with probability \(1\).

To achieve that, note that regardless of whether we chose to sample or not to sample an element in a given round, we modified the problem by reducing the size of the remaining stream by one, and if we sampled an element we reduced the number of yet-to-be-sampled elements by one as well. In a situation where only a single element is left to be sampled and we need another sample, we have a \(k=1\), \(N=1\) subproblem. This suggests the approach

from random import random

def sample(it, N, k):
    for e in it:
        if random() < k/N:
            yield e
            k -= 1
        N -= 1

Optionally, we could also short-circuit via something like if not k: return. A moment’s reflection shows that this algorithm always produces \(k\) elements and is correct.

This algorithm is Knuth’s Algorithm S, “selection sampling” (from The Art of Computer Programming, Vol. 2, §3.4.2). It is short, sweet, and “online” – we can yield instead of waiting for the stream to finish. It is the right approach for many applications, including handling classic magnetic tapes which are linear by nature.

Streams of unknown length

What should we do if we are not being given the length of the stream in advance? What can we do?

One obvious thing we cannot do anymore is to have an online solution that yields its results. Since the stream could go on for many more elements or stop right now and each element needs to have the same chance of making it, any solution will have to have the form

def sample(it, k):
    results = []

    for e in it:
        # do things and add to results
        ...

    return results

Heaps

One first idea might be to sort random numbers. We could for instance maintain a min-heap, along the following lines:

import heapq
from random import random

def sample(it, k):
    h = [(random(), next(it)) for _ in range(k)]
    heapq.heapify(h)  # Tuples compare by first value first.

    for e in it:
        # Push e on heap, then pop smallest.
        heapq.heappushpop(h, (random(), e))

    return [e for _, e in h]

Somewhat surprisingly, heapify is implementable in \(O(k)\). The use of heappushpop here makes sense since heap operations are typically “modify the heap as requested, then recreate the heap property”. Since here each push is followed by a pop, heappushpop allows us to skip one heap recreation. It runs in \(O(\log k)\) for a heap of size \(k\).

Overall this algorithm is \(O(k) + O(N \log k)\) in time.

There is an issue with this approach though: It relies on each random() number being truly different from each other one. Since Python’s floats don’t have infinite precision, collisions are possible and due to the birthday paradox are actually more likely than intuitively expected.⁴ If our streamed elements are not comparable, this will make the above algorithm crash at runtime – and if they are, uniformity will be violated. The only real solution to this I can see is to modify the comparison to have random tie breaks which breaks the deterministic runtime behavior and is not pleasant to implement.

Quickselect

While it doesn’t apply in the streaming case, it’s good to remember that quickselect is the standard algorithm for finding the \(k\)th smallest element in an unsorted list, or the \(k\) smallest elements. However since it requires modifying the full array, it isn’t the best approach for the problem at hand. (It also suffers from the same floating point collision issue.) It’s still very good to know quickselect, which is why I mention it here.

Reservoir sampling

The real beauty of the streaming version of this question comes out in this solution. Let’s go back to the template we know we will require:

def sample(it, k):
    results = []

    for e in it:
        # do things and add to results
        ...

    return results

We can view this question in several ways. One of them is to think of it as some sort of adversarial game: Let’s say you are the algorithm and the interviewer is the stream. The algorithm can do whatever it wants with its inputs, but the stream can decide to either give another element or end itself and ask you to “show your work” in the least convenient moment possible.

Thinking along these lines creates one realization: We have to be correct for all kinds of values of \(N\) and \(k\), including \(N = k\). In that case there isn’t much we can do, we need to return the first \(k\) inputs. So let’s code that:

def sample(it, k):
    results = [next(it) for _ in range(k)]

    for e in it:
        # do things and add to results
        ...

    return results

This now is at least correct for the \(N = k\) case. Can we make it correct for \(N = k + 1\) as well? (You see where this is going.) For \(N = k + 1\), what’s the chance of any given element to be in the results array? It will be \(\frac{k}{k+1}\). Essentially, if we see another element after the guaranteed first \(k\), what choice do we have? We can choose it or reject it with a certain probability, and if we choose it, we can evict a previously chosen element with some probability. Since the stream could end right after that additional element, its probability of being chosen has to be correct, namely k / N which is k / stream_len_so_far in this case.

This is the beauty of this algorithm: We have to assume a solution can be found, and under that assumption we only have that choice:

from random import random

def sample(it, k):
    results = [next(it) for _ in range(k)]

    for stream_len, e in enumerate(it, k + 1):  # Start with k + 1
        if random() < k / stream_len:
            # Choose `e`.
            ...

    return results

This is almost the full algorithm. The remaining question is what to evict? Given the uniform sample we want to produce, it stands to reason that we want to select elements for eviction from the pool results – let’s actually call it a reservoir – with uniform probability:

from random import random, randrange

def sample(it, k):
    reservoir = [next(it) for _ in range(k)]

    for stream_len, e in enumerate(it, k + 1):
        if random() < k / stream_len:
            reservoir[randrange(k)] = e

    return reservoir

A slight optimization requires only a single random choice here:

from random import randrange

def sample(it, k):
    reservoir = [next(it) for _ in range(k)]

    for stream_len, e in enumerate(it, k + 1):
        if (i := randrange(stream_len)) < k:
            reservoir[i] = e

    return reservoir

That is a neat algorithm. But is it correct? We arrived at it by assuming something like this is possible; in particular, we chose the probabilities such that the final element’s chance of ending in the chosen set is correct.

Given enough time, good candidates arrive at a short inductive proof of the correctness of this algorithm. For a given \(k\), the probability for any element to end up in the final set is obviously correct if \(N = k\). Assuming the algorithm is correct for a given stream length \(N\), moving to a stream length \(N + 1\) requires one more iteration. We already know the final new element is chosen with the right probability, but what happens to the other elements? Their previous chance to end up in reservoir was \(k / N\). Their chance to still be in the reservoir after one more iteration is that, times (1) the chance of the final element to be discarded, plus (2) the chance of the final element to be chosen, but another element to be evicted. This telescopes nicely since

So the probabilities of the previous elements are updated in the correct way too.

This algorithm is known as reservoir sampling. It’s a classic from the literature; Knuth calls it Algorithm R in TAOCP Volume 2.

Reservoir sampling with skipping ahead

Fancier versions of reservoir sampling exist. One interesting variant is if we assume we can cheaply skip ahead in the iterator. If the length of the stream N is large compared to the number of selected elements and there is processing work for each element that pays off to be avoided where possible, this can be a powerful idea.

Let’s invent an iterator method where say skip(it, j) skips and discards j elements from the iterator (making skip(it, 0) a no-op). The skipping version of reservoir sampling can then be written as

from math import exp, floor, log, log1p
from random import random, randrange

def sample(it, k, done=object()):
    reservoir = [next(it) for _ in range(k)]
    w = 1.0
    while True:
        w *= exp(log(random()) / k)
        skip(it, floor(log(random()) / log1p(-w)))
        if (e := next(it, done)) is done:
            break
        reservoir[randrange(k)] = e

    return reservoir

Here w is the probability that the next element of the stream will be accepted, and log1p(-w) computes \(\log(1 - w)\) without losing precision when w is small.⁵

This version of reservoir sampling was published in 1994 by Kim-Hung Li, see Reservoir-sampling algorithms of time complexity \(O(n(1 + \log(N/n)))\).

There are other interesting variants of reservoir sampling, e.g., sampling from non-uniform distributions. But let’s end with something a bit farther afield that deserves to be better known.

Online shuffling

Let’s say we have a large dataset (e.g., a table in SQL) which comes in the form

  # |   value
-------------
  1 |  entry1
  . | .
  . | .
  . | .
  N |  entryN

We now want to do “online shuffling”, i.e., go through the dataset in a random order, visiting each sample exactly once. We also want to be preemptible – our algorithm can be stopped and later restarted from a checkpoint. Perhaps we would even like this to be seekable: Quickly asking for the entry that comes at position i of the shuffle.

What we are asking for is equivalent to selecting one of the \(N!\) many permutations of \(N\) items, then doing lookups of the form

for i in range(N):
    sample = samples[p(i)]
    ...

This is similar to the sampling question in that doing this requires a scratch buffer of size \(O(N)\) to keep track of what has been sampled so far. For instance, this could be done like:

import numpy as np
from random import randrange

def permit(N):
    L = np.arange(N)
    for i in range(N):
        j = i + randrange(N - i)
        L[i], L[j] = L[j], L[i]
        yield L[i]

However, this is a heavy amount of state for our checkpoint: For a dataset of 1B entries we would use 8 GiB just for the np.arange(N). It’s also not quickly seekable in that the full array needs to be assembled before we can do lookups.

Block ciphers

If we want to deal with far less state (say, a few kilobytes), the problem cannot be solved exactly. But it can be solved approximately, and one can prove the approximation is good.

One trick is to use an idea from cryptography: A block cipher is a symmetric encryption algorithm which takes a secret key k and buffers of a fixed length and maps them to encrypted buffers of the same length:

plain text buffer: B
encrypted buffer:  encrypt(B, k)
decrypted buffer:  B == decrypt(encrypt(B, k), k)

Both the unencrypted input B and the encrypted buffer encrypt(B, k) have the same number of bits, say \(n\). In particular, this means that for a given key k the functions

are bijections (one-to-one) and each is the inverse of the other. Here, \(\set{0, 1}^{n}\) is the set of all buffers of 0s and 1s of length \(n\), which is a set of \(2^n\) elements.

Pseudorandom permutations

How are we going to use this block cipher? We will use \(n = 64\) and interpret a given index i of type uint64 as a block of 64 bits and use a block cipher for that block size. We then encrypt that index to find where it maps to:

import numpy as np

def perm(i):
    B = np.array(i, dtype=np.uint64)
    k = KEY  # Global in this example.
    return encrypt(B, k).item()

def permit(N):
    for i in range(N):
        yield perm(i)

Since \(\mathtt{encrypt}(\dotid, \mathtt{k})\) is one-to-one, each possible bit pattern is produced exactly once. perm is an actual permutation of the set of uint64s (a set of \(2^{64}\) elements).

This works as described if the (implausibly large) table has size exactly \(N = 2^n = 2^{64}\). Realistically, it’s smaller and we need some additional ideas.

Let’s say we have a block cipher for any power of two and a table of any size. We can use as block size \(n\) the smallest number such that \(2^n \ge\) the table size. Worst case, this additional padding at the end is roughly half the range of the permutation. When we apply perm and hit the padded region, we will simply iterate perm until we get out of it – since the permutation is one-to-one, that will eventually happen, and since the padded part is at most half, it happens with probability at least 0.5 every iteration. This trick is known as cycle walking.

Is 0% enough?

Note that the permutation is parameterized only by the key k. If k has m bits, that means we can only produce \(2^m\) different permutations, which is vastly fewer than the full number of permutations of the set of uint64s, which is \((2^{64})!\), a number that has about \(3.47 \times 10^{20}\) digits. By contrast, the number of atoms in the observable universe has about 80 digits! Combinatorics creates large numbers very fast.⁶ The actual fraction of permutations reachable by this scheme is very close to \(0\%\).

So we will only ever get an absurdly tiny fraction of all possible permutations this way. How far away from truly random permutations are we with these pseudorandom permutations?

The answer depends on the quality of the encrypt function. In many scenarios, cryptographers could prove that one cannot actually tell the difference! To learn more, check out Feistel networks⁷ and a result known as the Luby-Rackoff theorem.⁸

Cheap parameterized permutations

On the other end of the spectrum, computationally cheap parameterized permutations come from mixers (or finalizers) of hash functions like MurmurHash3. A very influential 2011 blog post by David Stafford proposed using

uint64_t mix13(uint64_t x)
{
    x ^= x >> 30;
    x *= 0xBF58476D1CE4E5B9UL;
    x ^= x >> 27;
    x *= 0x94D049BB133111EBUL;
    x ^= x >> 31;
    return x;
}

MurmurHash3 and this blog post were the starting point of all kinds of results around fast permutations. Stafford’s suggestion made it into JDK 8 in 2013 and was also cited as canonical in Fast Splittable Pseudorandom Number Generators by Steele, Lea and Flood in 2014. Since 2018, Pelle Evensen has had several blog posts around improvements. These fast permutations satisfy certain statistical properties but not others (e.g., mix13 above sends an all zero buffer to zero; this can be fixed with extra xors). They are certainly not “cryptographically strong”, but their statistical properties can be surprisingly good and they are extremely cheap.

Much more can be said about random sampling and shuffling. But this blog post is long enough and making it longer will not illustrate the point any better.