Efficiently shuffling items in $N$ buckets using $O(N)$ space

Question

I’ve run into a challenging algorithm puzzle while trying to generate a large amount of test data. The problem is as follows:

• We have $N$ buckets, $B_1$ through $B_N$. Each bucket $B_i$ maps to a unique item $a_i$ and a count $k_i$. Altogether, the collection holds $T=\sum_1^N{k_i}$ items. This is a more compact representation of a vector of $T$ items where each $a_i$ is repeated $k_i$ times.

• We want to output a shuffled list of the $T$ items, all permutations equally probable, using only $O(N)$ space and minimal time complexity. (Assume a perfect RNG.)

• $N$ is fairly large and $T$ is much larger; 5,000 and 5,000,000 in the problem that led me to this investigation.

(EDIT: further research instigated by @YuvalFilmus’s comment shows that this is equivalent to weighted sampling without replacement, a search term that leads to quite a lot of research.)

Now clearly the time complexity is at least $O(T)$ since we have to output that many items. But how closely can we approach that lower bound? Some algorithms:

• Algorithm 1: Expand the buckets into a vector of $T$ items and use Fisher-Yates. This uses $O(T)$ time, but also $O(T)$ space, which we want to avoid.

• Algorithm 2: For each step, choose a random number $R$ from $[0,T-1]$. Traverse the buckets, subtracting $k_i$ from $R$ each time, until $R<0$; then output $i$ and decrement $k_i$ and $T$. This seems correct and does not use extra space. However, it takes $O(NT)$ time, which is quite slow when $N$ is large.

• Algorithm 3: Convert the vector of buckets into a balanced binary tree with buckets at the leaf nodes; the depth should be close to $\log_2{N}$. Each node stores the total count of all the buckets under it. To shuffle, choose a random number $R$ from $[0,T-1]$, then descend into the tree accordingly, decrementing each node count as we go; when descending to the right, reduce $R$ by the left count. When we reach a leaf node, output its value. It uses $O(N)$ space and $O(T\log{N})$ time.

• Algorithm 3a: Same as Algorithm 3, but with a Huffman tree; this should be faster if the $k_i$ values vary widely, since the most often visited nodes will be closer to the root. The performance is more difficult to assess, but looks like it would vary from $O(T)$ to $O(T\log{N})$ depending on the distribution of $k_i$.

Algorithm 3 is the best I’ve come up with. Here are some illustrations to clarify it:

Does anyone know of a more efficient algorithm? I tried searching with various terms but could not find any discussion of this particular task.

Answer 1

Step 1. Store the buckets cumulatively. E.g. instead of storing $[3, 5, 2, 2]$ you store $[3, 8, 10, 12]$. This allows you to find using binary search in $\log n$ time what bucket the $k$th (out of $T$) item belongs to (and by extension which $a_i$ to output).

Step 2. Use a memory-less permutation of the range $[1, T]$. I personally like Sometimes-Recurse Shuffle. A memory-less permutation can give you $\pi(k)$ directly for arbitrary $k$ in $\log T$ time. Some example code I have written a couple years back show how simple it is (with toy PRF implementation - replace that): https://gist.github.com/orlp/33535eefce782a59e185e4a971cda1a3.

Step 3. To actually shuffle simply enumerate $k = 1, 2, \dots$, compute $\pi(k)$ using the memory-less permutation, and find which bucket $\pi(k)$ belongs to output the proper $a_i$. Total time complexity is $\log(T)$ per item, or $T \log(T)$ for the whole permutation.

The advantage is that you don't have to compute the whole permutation, you can also just directly compute the $k$th element in the permutation.

Answer 2

Here's a pretty simple algorithm for weighted random selection without replacement, which is O(N) space and stochastic O(1) time per element selected (i.e. O(T) time for a complete permutation).

It's a combination of two common algorithms for weighted random selection: the rejection method (I can't find a reference for this one; it's as old as the hills although it frequently gets rediscovered, as for example the frequently-cited Lipowski & Lipowska 2011) and the alias method (Walker 1977; Kronmal & Peterson 1979). A good discussion of both of these methods (and many more) can be found in Chapter 3 of Non-Uniform Random Variate Generation (Luc Devroye, 1986) which is freely-downloadable through that link. (And kudos to Devroye for his rant about academic bloodsuckerspublishers. I can't find any unpaywalled copies of the 40-year-old papers cited above, so Devroye is the best reference I have. The algorithms are described below.)

The basic problem with both of these methods is that they require O(N) setup time, and the fact that the weight distribution varies after each selection is extracted would require reprocessing at every step, resulting in O(NT) time. In addition, the rejection method only provides O(1) selection time over a restricted range of possible weight distributions [Note 1], and we cannot guarantee that requirement for all distributions which might occur during a permutation. However, putting the two algorithms together solves both problems:

1) We use the alias method (slightly modified) to guarantee that the rejection method's distribution requirement is honoured; and

2) Because of the flexibility provided by the rejection method, we do the O(N) auxiliary processing less then once every N selections, so that it has an O(1) amortised contribution to the cost of doing a selection.

The rejection method

The rejection method is incredibly simple. We start with some weight $\omega_\alpha$ which is at least as great as any weight in the distribution. (That's the required precomputation; clearly, we could compute $\omega_\alpha = max_{i=1}^N(\omega_i)$

At each step:

• Repeatedly generate two random variates $u, v$ over the uniform ranges $[1, N]$ and $[1, \omega_\alpha]$ respectively [Note 2] until $v < \omega_u$, and then return $u$.

In effect, this corresponds to firing a dart at an $N \times \omega_\alpha$ grid, in which the bottom $\omega_i$ squares in column $i$ are coloured. If the dart falls on a coloured square, we return the number of that column; otherwise, we try again. It should be clear both that the values returned obey the desired probability distribution, and that the probability of returning a value on any given iteration is $\sum \omega_i \over { N \omega_\alpha }$.

If we knew that ${\omega_\alpha \over \omega_{mean}} \lt c$ for some constant $c$, then the expected number of iterations would be $c$. (Or, to put it another way, the expected number of rejections would be $c-1$.)

But, of course, without doing some more work we cannot make this assertion. However, it's worth noting that if $\omega_{mean}$ is not too small, then removing an element has little impact on the ratio. Consequently, we can avoid recomputing $\omega_\alpha$ on every step. But that's getting a bit ahead of ourselves.

The alias method

Suppose we take the chessboard described above and attempt to level it out by filling in small columns using coloured squares from the big columns. It turns out that we can always find an arrangement in which there are only two colours in any column, and in which all of the columns are the same height (within one square, since there is no guarantee that the total weight evenly divides N). We call the alternative value for each column it's alias ($a_i$) and we store two weights for each column: $\omega^0_i$, the number of squares corresponding to the value at index $i$, and $\omega^1_i$, the total number of coloured squares in that column. [Note 3]

Evidently, we can choose $\omega_\alpha = \lceil T / N \rceil$ and then run a slightly modified version of the rejection algorithm. The only difference is that the iteration is changed to:

• Repeatedly generate two random variates $u, v$ over the uniform ranges $[1, N]$ and $[1, \omega_\alpha]$ respectively [Note 2] until $v < \omega^1_u$, and then select

$$\begin{cases} u, & \text{if } v \lt \omega^0_u \\ a_u, & \text{otherwise} \end{cases} $$

The alias table (in this version) requires storing $3N$ numbers, all of them in the same range as the problem's inputs. That means that when we have reduced $T$ to $3N$, we can switch to an alternative algorithm without using any more space: collect all the remaining values in a single array of size $3N$, shuffle it randomly, and then generate its elements in sequence. That lets us assume that $\omega_\alpha$ is at least 3, which is convenient because it lets us limit the expected number of rejections without recomputing the alias table too often.

As noted above, if we recompute the alias table no more often than every $N$ steps, we will contribute amortised O(1) to the cost of a single step. So a simple approach is to recompute the alias table every time $T$ is an exact multiple of $N$, until $T = 3N$. In that case, it will always be the case that

$${\omega_\alpha \over \omega_{mean}} \lt { 4 \over 3}$$

so the expected number of rejections is (much) less than $1/3$.

Indeed, we might be prepared to live with a larger number of rejections. If, for example, we wanted to limit the expected number of rejections to 1, it would suffice to recompute the alias table every time $T$ fell to half its previous value (again, stopping at $T = 3N$), which would mean that we would do $lg N$ recomputations. However, experimentation indicates that the cost of recomputing the alias table is so small that this is not worthwhile.

---

Notes

1. The algorithm is O(1) for any given weight distribution, although the constant can be quite large, but if we do the analysis over all possible distributions which sum to $T$, the worst-case expected time increases without bound as $|T|\rightarrow \infty$. In practice, we will want to limit the size of the constant.

2. In practice, we normally generate a single variate over the range $[0, N\omega_\alpha)$ and then extract the two required variates through division and modulo by $N$. I've written this answer using the usual mathematical convention of 1-indexing lists, but clearly the arithmetic will be tidier when implemented using a 0-indexed programming language.

3. In Walker's paper, and indeed in most descriptions of the alias method, it is assumed that the weights are infinite-precision real values, in which case all of the columns will be exactly the same height. Of course, that is not true of any practical implementation; usually, the weights are adjusted very slightly to even out the columns. In the case that the weights are integers, this adjustment is not possible, so the second weight must be stored, although only a single bit is required since each column's total weight is either $\lceil T / N \rceil$ and $\lfloor T / N\rfloor$. We don't even have to store this bit, if we ensure that the short columns are all at the end (which is easy to do); we could just store the number of short columns.

   However, here we are going to want to modify the weights dynamically (by subtracting one each time an item is selected) so the column sizes will vary unpredictably as the algorithm proceeds. So we store three values instead of two.