This is a blog 3 of a series of blogs about algorithmic challenges I came across when creating tgen.

In this blog we will tackle:

1. Create a data structure with $$$\mathcal O(n)$$$ time construction that can return $$$i \in [0, n)$$$ with probability proportional to $$$w[i] \geq 0$$$ in $$$\mathcal O(1)$$$ time.

It is trivial to create such a data structure with $$$\Theta(\log n)$$$ query time: with a prefix sum array $$$p$$$ with values in $$$[0, S]$$$, where $$$S = \sum_i w[i]$$$, just pick a uniformly random value $$$x \in [0, S)$$$, and then do a binary search to find and return the first index $$$i$$$ such that $$$p[i] \gt x$$$.

Turns out that we can improve the query to $$$\mathcal O(1)$$$ time, with linear construction. This is the alias method.

If all weights were the same, $$$S / n$$$, then it would be trivial: just return next(0, n-1). If not, some indices are small ($$$w[i] \lt S / n$$$), and some are big ($$$w[i] \geq S / n$$$). The idea is to "transfer" the exceeding probability of the small indices to the big ones.

To do that, after we choose a random index i = next(0, n-1), we will either return i with probability mass[i] / S, or return alias[i]. Note that there is no chaining, that is, we don't follow many alias pointers; we either return i or alias[i]. Now we are left to show that it is always possible to choose an assignment for alias[i] and mass[i] in such a way that the probabilities are maintained.

To start, we scale every weight by $$$n$$$ to make everything an integer, so outcome $$$i$$$ has mass $$$n \cdot w[i]$$$. The average scaled mass is exactly $$$S$$$. Maintain two queues: $$$\mathrm{small}$$$ for indices with current mass $$$ \lt S$$$, and $$$\mathrm{big}$$$ for indices with current mass $$$\geq S$$$. While both are nonempty, pop $$$s$$$ from $$$\mathrm{small}$$$ and $$$b$$$ from $$$\mathrm{big}$$$. Index $$$s$$$ is finalized: we keep $$$\mathrm{mass}[s]$$$ as the threshold for returning $$$s$$$, and set $$$\mathrm{alias}[s] = b$$$, so the leftover $$$S - \mathrm{mass}[s]$$$ of the pick is given to $$$b$$$. This leftover is at most $$$S$$$ (since $$$\mathrm{mass}[s] \geq 0$$$), and $$$\mathrm{mass}[b] \geq S$$$, so we never subtract more from $$$b$$$ than it has. After subtracting that leftover from $$$\mathrm{mass}[b]$$$, we put $$$b$$$ back into the appropriate queue (depending on whether its remaining mass is small or big). Each step permanently finalizes one index.

At the end, all indices remaining in $$$\mathrm{big}$$$ must have mass exactly $$$S$$$ (because the existence of an index with mass $$$ \gt S$$$ implies another with mass $$$ \lt S$$$), so those indices are finalized as a single outcome (they are their own alias).

struct alias_method {
	int n;
	long long S;
	std::vector<long long> mass;
	std::vector<int> alias;

	alias_method(const std::vector<int>& w) : n(w.size()), alias(n) {
		S = 0;
		for (auto x : w) S += x;

		std::queue<int> small, big;
		for (int i = 0; i < n; i++) {
			mass.push_back(n * w[i]);
			if (mass[i] < S) small.push(i);
			else big.push(i);
		}

		while (!small.empty() and !big.empty()) {
			int s = small.front(); small.pop();
			int b = big.front(); big.pop();

			alias[s] = b;
			mass[b] -= S - mass[s];

			if (mass[b] < S) small.push(b);
			else big.push(b);
		}

		while (!big.empty()) {
			int b = big.front(); big.pop();
			alias[b] = b;
			mass[b] = S;
		}
	}
};

Algorithm 1: alias method construction.

To sample, pick a uniform index and compare a uniform sample in $$$[0, S)$$$ against the threshold stored in that index:

int next(const alias_method& a) {
	int i = next(0, a.n - 1);
	long long x = next(0, a.S - 1);
	return x < a.mass[i] ? i : a.alias[i];
}

Algorithm 2: alias method query.

Theorem 1: Algorithm 2 returns index $$$i$$$ with probability $$$w[i] / S$$$.

Proof: Algorithm 2 picks an index $$$j$$$ uniformly at random, then returns $$$j$$$ if $$$x \lt \mathrm{mass}[j]$$$ and $$$\mathrm{alias}[j]$$$ otherwise, with $$$x$$$ uniform in $$$[0, S)$$$. So, conditioned on picking $$$j$$$, we return $$$j$$$ with probability $$$\mathrm{mass}[j]/S$$$ and $$$\mathrm{alias}[j]$$$ with probability $$$1 - \mathrm{mass}[j]/S$$$.

Initially $$$\mathrm{mass}[i] = n \cdot w[i]$$$. Each pairing of a small index $$$s$$$ with a big index $$$b$$$ keeps $$$\mathrm{mass}[s]$$$ units with $$$s$$$ and moves $$$S - \mathrm{mass}[s]$$$ units from $$$b$$$ to the alias of $$$s$$$.

The probability of returning $$$i$$$ is therefore

where $$$A_i = \mathrm{mass}[i] + \sum_{\mathrm{alias}[j] = i} (S - \mathrm{mass}[j])$$$. The first term is the case where we pick $$$i$$$ and keep it; each term in the sum is the case where we pick some $$$j$$$ and take the alias $$$i$$$.

Initially $$$A_i = n \cdot w[i]$$$, since there are no aliases yet. A pairing does not change $$$A_s$$$ or $$$A_b$$$. We do not change $$$\mathrm{mass}[s]$$$, so $$$A_s$$$ stays the same. For $$$A_b$$$, we subtract $$$S - \mathrm{mass}[s]$$$ from $$$\mathrm{mass}[b]$$$, but we also set $$$\mathrm{alias}[s] = b$$$, which adds the same $$$S - \mathrm{mass}[s]$$$ to the sum in $$$A_b$$$. Thus $$$A_i = n \cdot w[i]$$$ at the end, so we get the final probability of returning $$$i$$$:

$$$\square$$$

References

Alias method (Wikipedia): https://en.wikipedia.org/wiki/Alias_method

Keith Schwarz's exposition: https://www.keithschwarz.com/darts-dice-coins/