On shuffling

One of the most fundamental operations of randomized algorithms is shuffling: Given an array of \(n\) items, randomly permute their order, such that every possible order (every permutation) is equally likely. The classic algorithm for this is due to Fisher and Yates. This post is about shuffling and a neat related puzzle which occurred to me about a year ago.

Fisher-Yates shuffling

Suppose you had a “repository” of remaining items to be shuffled. In each step, you pick one of these items uniformly at random, removing it from the repository, and adding it to a list of shuffled items.

The natural data structure for such a repository is an array (or more precisely, a “slice”): By choosing an index, you can efficiently choose an item uniformly at random.

Because we do not care about the order of items in the repository, we can efficiently remove an item from the repository by swapping it with the last item, and then “removing” that item by restricting the slice to the remaining range of slots.

function shuffled(items)
	local n = #items
	local repo = {}
	for i = 1, n do
		repo[i] = items[i] -- copy the list
	end
	local shuf = {}
	for i = 1, n do
		local j = math.random(1, #repo)
		shuf[i] = repo[j]
		repo[j] = repo[#repo]
		repo[#repo] = nil -- shrink repo
	end
	return shuf
end

(The examples here use Lua, but translate straightforwardly to other programming languages.)

This is however a bit more complicated than necessary. It is also wasteful, constructing two additional tables. More efficiently, we can keep the shuffled result and the repository of remaining items in the same array as we go: After the i-th loop iteration, the slice items[1], …, items[i] is a random, shuffled subset of the items; the slice items[i+1], …, items[n] is the repository of remaining items. We maintain this invariant by picking a random item from the repository to swap to the i-th place in the i-th iteration:

function shuffle(items)
	local n = #items
	for i = 1, n - 1 do -- last iteration changes nothing, can be skipped
		local j = math.random(i, n) -- choose from repository
		-- Swap chosen item from repository to shuffled slice.
		items[i], items[j] = items[j], items[i]
	end
end

This is the classic in-place Fisher-Yates shuffle. (To get a shuffled copy instead, you would simply copy the list of items first.)

Fun fact: It seems to happen not too rarely that programmers implement “shuffling by random swaps”, but don’t restrict the swaps properly (e.g., they write a loop similar to the above, but pick j = math.random(1, n)), which leads to a very skewed distribution where some permutations are much more likely than others (you can either try this out by writing some tests, or by computing the probabilities for some small examples by hand).

Clearly, at least linearly many swaps are necessary for every permutation to be possible at all. And because at least half of all permutations require at least linearly many swaps, linearly many swaps are even necessary in expectation. In this sense, the linear runtime of Fisher-Yates is optimal; we can not hope to find an algorithm with sublinear expected runtime here.

However, we often only wish to sample a subset of a list, the size of which may not be known in advance. Can this be done more efficiently than shuffling the entire list? The answer is yes.

The ballot box

Say you want to draw some distinct items from a set, with each subset (and each order) being equally likely – but there are potentially many candidates and you do not know how many items you will want to draw ahead of time, so just shuffling all items might end up being rather wasteful if you only end up drawing a few.

A “ballot box” only supports a single operation: extracting an item chosen uniformly at random (“drawing” an item). We can treat a ballot box as a function (a closure) which upon each call returns the drawn item: draw = ballot_box(n) gives you a ballot box containing the indices \(1, 2, ..., n\), which upon each call index = draw() returns and removes one of the remaining indices in the box chosen uniformly at random.

To shuffle a list, you could simply call draw() repeatedly, recording all the picks in an array:

function shuffled(list)
	local draw = ballot_box(#list)
	local shuf = {}
	for i = 1, #list do
		shuf[i] = list[draw()]
	end
	return shuf
end

We make a first attempt at implementing such a “ballot box”:

Rejection sampling

We can record previously chosen indices in a hash set; we can then repeatedly choose a random index from \(1\) to \(n\), rejecting any indices that were recorded previously, until we get a fresh index:

function ballot_box(n)
	local chosen = {}
	return function()
		local i
		repeat
			i = math.random(1, n)
		until not chosen[i]
		chosen[i] = true
		return i
	end
end

Strictly speaking, this algorithm is not guaranteed to terminate at all once the first choice has been made: We could get perpetually unlucky, such that the repeat loop would never be left. But the chance for this is virtually zero as we are about to see. Instead of a “true” worst-case runtime, we thus speak of an expected runtime, treating the runtime as a random variable depending on the random numbers math.random gives us. ¹

What is the expected asymptotic runtime if we call draw \(n\) times? Perhaps surprisingly, the answer is \(O(n \log n)\).

To see why, we look at the expected runtime of a single draw.

We assume that each line in the pseudocode shown above runs in expected amortized constant time. In particular, we assume that the hash map access is expected constant time. This is reasonable as keys are drawn uniformly at random from the integers \(1, ..., n\).

It is also possible to draw uniformly distributed random integers in a certain range in expected constant time. ²

Thus it suffices to count the number of steps until we obtain a fresh index. This is a random variable following a geometric distribution.

Let \(m\) be the number of indices that have already been drawn. Let \(p = \frac{m}{n}\) be the probability that a random choice is a “miss”, meaning we draw an index we have already seen previously. The probability that we need exactly \(k\) steps is \((1-p) p^k\), decreasing exponentially in \(k\); the expected number of steps is \(\frac{1}{1 - p} = \frac{n}{n - m}\).

This will be very efficient in the beginning, when few items have been chosen; it will get less and less efficient the more often draw() is called, the more likely it becomes that you have to reject an index. This gets worst for the last draw, where you only have a \(\frac{1}{n}\) chance of hitting the single remaining index, and thus can expect linearly many tries.

However, for a sequence of \(k\) draws starting from a fresh ballot box, we have \(k \cdot \sum_{m=0}^{k-1} \frac{1}{n - m} = k \cdot \sum_{i=1}^k \frac{1}{i} = k \cdot H_k = O(k \log k)\) steps in expectation, where we use the fact that \(H_k = O(\log k)\) ³. If we blame each of the \(j\) operations for an equal part of this runtime, we see that this cute approach is not too bad: in expectation, each draw() operation has amortized logarithmic runtime.

This kind of problem is also known as the Coupon collector’s problem: There are \(n\) different coupons. You want to collect them all. Every time you go buy something, you get a coupon, chosen uniformly at random. How long does it take until you have all \(n\) different coupons? By an analogeous argument, the answer is \(O(n \log n)\).

But we want to do better. We want asymptotically optimal amortized, expected constant time complexity per draw(). How can you achieve this?

You might want to think about it yourself for a bit.

function ballot_box(n)
	local n_chosen = 0
	local chosen = {}
	local stack

	return function()
		-- Second phase: Take from stack of shuffled items.
		if stack then
			return table.remove(stack)
		end

		-- First phase: Use rejection sampling.
		-- This is expected constant time because there's a probability of at least 1/2
		-- of finding an index that has not been chosen before.
		local i
		repeat
			i = math.random(1, n)
		until not chosen[i]
		chosen[i] = true

		-- Transition to second phase once some fixed fraction of items have been chosen.
		-- The linearly many choices we have done so far pay for this.
		n_chosen = n_chosen + 1
		if n_chosen > n/2 then
			stack = {}
			for i = 1, n do
				if not chosen[i] then
					table.insert(stack, i)
				end
			end
			shuffle(stack)
		end

		return i
	end
end

function ballot_box(n)
	local perm = {}
	local i = n
	return function()
		local j = math.random(i)
		local index = perm[j] or j
		perm[j] = perm[i] or i
		perm[i] = nil
		i = i - 1
		return index
	end
end

function ballot_box(n)
	local n_chosen = 0
	local n_choices = n
	local choices = {}

	local function shrink()
		local new_choices = {}
		for i = 1, n_choices do
			if choices[i] ~= false then
				table.insert(new_choices, choices[i] or i)
			end
		end
		n_choices = n_choices - n_chosen
		n_chosen = 0
		choices = new_choices
	end
	
	return function()
		local i
		repeat
			i = math.random(1, n_choices)
		until choices[i] ~= false
		local choice = choices[i] or i
		choices[i] = false
		n_chosen = n_chosen + 1
		if n_chosen >= n_choices/2 then
			shrink()
		end
		return choice
	end
end