# Random sampling data structure with removal

I'd like a data structure with the following operations:

1. create a new instances from an array of floating point weights.
2. randomly sample, returning an item with probability proportionate to its weight.
3. set the weight of a specific item to 0.

Without the third option this is pretty standard: Just implement an alias method sampler using Vose's algorithm, but the ability to do a dynamic update makes it harder.

The first operation is intrinsically O(n) (or at least it is if you don't make the second operation O(n)), but it would be nice to be able to do the other two in amortized O(1), or at least O(log(n)). Sampling will be significantly the more common operation, so ideally I would like to have it not be much slower (either in complexity or constant factors) than using the alias method.

One option is to just rebuild the sampler every time you do a remove, but that's O(n) which is less than ideal.

The following is currently my best bet, which builds on that to amortise it a bit:

We keep a copy of the set of weights, and build an alias method sampler initially. When we set the weight of an item to zero we update our weights table appropriately. Then when we sample from our alias table, we check the weight. If it's zero then we rebuild the alias table with the updated weights and draw again.

This should work OK for amortising some of the cost, but unfortunately the workload I'm likely to want to put it through is probably pessimal for it: Items that are going to be removed will have been drawn from the sampler, so will tend to be of high weight. This means that it will end up hitting the rebuild case fairly commonly.

So I'm hoping for something better. Any suggestions?

• Memory allocators are designed to solve a similar problem: blocks of various sizes can be allocated and freed, while making an attempt to keep them packed as tightly as possible to minimize total contiguous memory usage. If for each weight we allocate a block of that size and fill it with copies of its index, this "data structure" can be sampled by picking a random address and reading it until we find one that contains a valid index. Not an answer but maybe an interesting connection. – Dan Brumleve Feb 27 '17 at 17:27
• Hmm. The problem with that is that memory allocators also get to use O(weight) amount of space, which is somewhat impractical. – DRMacIver Feb 28 '17 at 9:15
• I've implemented Dan Brumleve's suggested, which is available as weightedDict.py at: github.com/google/weighted-dict – Marctar Jun 29 '18 at 14:06

Copying my comment on that from here:

There exist published algorithms that support sampling from discrete probability distributions in O(1) time, AND modifying the distribution in O(1) time per update:

These papers are a little unsatisfactory in that they are theoretical algorithms in a model that doesn't fully deal with issues of numerical precision and bit complexity. This can be an issue in practice if the probability distribution has some entries that are much larger than others. Consequently, in the application I had, I used some of the ideas from these papers, but implemented something else that was more practical.

EDIT: What I implemented was a bucketing system. Assume without loss of generality that each weight (non-normalized probability) is a power of 2.(Normalize so each weight is a 64-bit integer. Round each weight up to its nearest power of 2, then sample from the rounded weights instead. When you sample a value $i$ that has been rounded up to $2^k$, only actually take it with probability $i/2^k$, and otherwise resample. You get at most 1 repeat in expectation.) Keep all rounded weights with the same value $2^k$ in a single bucket indexed by $k$. As you do updates, keep track of the currently largest rounded weight $2^K$, the sum of the weights in each bucket, and the total weight $W$.

To sample, draw a random integer $x$ in $[1,W]$, start from the non-empty bucket that holds the largest weight $2^K$. If $x$ is at most the sum of the weights in the bucket, take the $j$th item in this bucket, where $j$ is something like $\lceil x/2^K\rceil$. Oherwise subtract the sum of the weights in the bucket from $x$ then recurse with the bucket that holds $2^{K-1}$.

If you have a weight distribution with very high ratio of largest weight to smallest weight, classify a weight as small if its ratio to the maximum weight is less than, say, $1/n^3$. Let $2^k$ be the smallest power of 2 larger than all small weights, and round each small weight up to $2^k$ (instead of its next largest power of 2). (This may be necessary in order to use 64-bit integer arithmetic to keep track of the exact total weight $W$.) If you happen to sample a small weight $i$ that has been rounded up to $2^k$, as before you only keep it with probability $i/2^k$, and otherwise resample. The probability of keeping it can be small, but the probability that you draw any such item is, say, $O(1/n^3)$, so you still do $O(1)$ resamples in expectation.

Now you have at most $3\log_2 n$ buckets, and worst-case (expected) time per sample is $O(\log n)$. As you delete weights, when the total weight of the small items reaches, say, a $1/n^2$ fraction of the total weight, reclassify the small items, taking some out of the "small-item" bucket. This happens $O(\log X / \log n)$ times, where $X$ is the ratio of largest weight to smallest weight.

You can get time $\Theta(\log n)$ for some distributions (e.g., for each $k\le \log n$ you have $2^k$ items of weight $n/2^k$). But in practice I was getting small constant sample time and update time.

• These are really interesting papers, thanks. I suspect I probably won't end up using them due to complexity, but I'm definitely going to read the papers in depth! – DRMacIver Feb 28 '17 at 9:16
• I edited the entry to describe the scheme I used. – Neal Young Feb 28 '17 at 12:21

Here's an idea:

Store the weights in the leaves of a balanced binary tree, and at each intermediate node store the sum of weights at the leaves underneath. Now we can sample it in log time by taking a weighted random walk from the root down to a leaf, using the stored sums to determine the probability of going left or right. And we can update a weight in log time by climbing up to the root from it updating the sum stored at each intermediate node.

• IMO this solution is easier to implement using a binary segment tree / Fenwick tree. – daniello Feb 27 '17 at 18:37
• Hmm. This is a nice solution, but I feel like the slow down on sampling might be unacceptable. You end up doing log(n) full draws from PRNG on each sample, which isn't that cheap. I'll definitely give this answer a try though! – DRMacIver Feb 28 '17 at 8:30
• You can use a single random draw x for each sample, if the draw represents, say, a uniform real in [0,1]. Draw x, then go find the first leaf L such that the cumulative probability of L and leaves before it is at least x. You still need $\log n$ time to find the leaf. – Neal Young Feb 28 '17 at 11:26