Sampling numbers from a weighted set that sum to constant value

Question

So I have a multi-set of positive integers $S = \{n_1, n_2, \dots\}$ with associated weights $W = \{w_1, w_2, \dots\}$. I want to sample some numbers, without replacement, from $S$ according to weights $W$ such that the sampled numbers sum to a given constant positive integer $k$.

The naive way to do it, as I understand it, is to find out all combinations of numbers from $S$ that sum to $k$ and add their weights and choose one of the combinations according to the sum of weights. But doing that is $O(|S|^k)$, which is unacceptably large.

What I'm doing right now is do a weighted shuffle of $S$ according to weights $W$, and picking number in the obtained order until the requirement is fulfilled. If, for example, the order obtained is $3, 2, 3, 2, 1$ and $k = 4$, then I'll pick $3$, ignore $2$ (as picking it will exceed $k = 4$) and pick $1$. This approach turns out to be $O(n\log\ n)$.

Does anyone have a better idea for doing this sampling in a way that is faster than exponential but more robust than what I'm currently doing now?

Answer 1

There's a close connection between counting the number of solutions and randomly sampling from the set of solutions. Any time you need to randomly sample, it's often helpful to ask yourself how you'd count the number of solutions, and then you can often turn that into a way to randomly sample.

So, one approach is to use dynamic programming to count the number ways to select numbers from $S$ that sum to $k$ (weighted by their weights $W$), then use that to help you sample from this space. Let me spell out the details more.

Define

$$f(S,W,k) = \sum_I \prod_{i \in I} w_i,$$

where $I$ ranges over all sets of indices such that $\sum_{i \in I} n_i = k$. Notice that if all the weights were 1, then $f(S,W,k)$ would count the number of ways to select these numbers; in general, with arbitrary weights, you can think of $f(S,W,k)$ as a weighted count, where you sum up the weights of each candidate combination, and the weight of a combination is the product of the weights of the numbers selected. You can compute $f(S,W,k)$ using dynamic programming, using the recurrence

$$f(\{n_1,\dots,n_j\},W,k) = f(\{n_1,\dots,n_{j-1}\},W,k) + w_j f(\{n_1,\dots,n_{j-1}\},W,k-n_j).$$

Now you want to sample a set $I$ with probability proportional to $\prod_{i \in I} w_i$. This can be done, using your algorithm for computing $f(S,W,k)$. In particular, if $S=\{n_1,\dots,n_m\}$, then flip a coin with heads probability

$${f(\{n_1,\dots,n_{m-1}\},W,k) \over f(\{n_1,\dots,n_m\},W,k)}.$$

If it is heads, then don't include $n_m$ in the solution: instead, recursively sample some numbers from $\{n_1,\dots,n_{m-1}\}$ that sum to $k$, and output that as the random sample. If it is tails, then do include $n_m$ in the solution: recursively sample some numbers from $\{n_1,\dots,n_{m-1}\}$ that sum to $k-n_m$, then add $n_m$ to that combination, and output that as the random sample. You can see that this induces the correct probability distribution on samples.

Overall, the running time will be comparable to the time to compute $f(S,W,k)$, or in other words, the running time will be $O(|S| \cdot k)$. This is much better than the $O(|S|^k)$ naive solution. It can still be very slow if $k$ is enormous, but if $k$ is not too big, it might be perfectly satisfactory.