Weighted sample of ~k elements from array in O(n) time?

Question

I have an array $a$ with $n$ elements, all of which have an associated weight. For example:

$a = \{ (A,2), (B,5), (C,9), ..., (Z,1) \}$, such that element $A$ has weight $w_A=2$, element $B$ has weight $w_B=5$ and so on. The array is not necessarily sorted.

Now, I would like to obtain a sample from this array where each element is selected without repetition and the probability is proportional to the weight of the element.

My goal is to sample approximately $k < n$ elements. I am emphasizing approximately because it does not have to be exactly $k$ but something close to $k$ is good enough. More precisely, I don't care if there is a deviation in the size of the final sample of $\pm \delta$ so long as $\delta$ is not too large in comparison to $k$ or $n$.

I know how to do the sampling in time $O(n + k \log n)$ by computing the cumulative weight $W$, then generating a random value in $[0,W]$ followed by a binary search to check which element this corresponds to.

However, my goal here is to understand whether it is possible to make a sampling in $O(n)$ by defining a probability $p_i$ for each element in such a way that I can decide independently for each element whether it should belong to the sample. Meanwhile, in the end I would still have a sample with approximately the size $k$ that I am looking for. If this is possible, then I could make a single pass in array $a$ after defining the probability $p_i$ accordingly which would require the procedure to define $p_i$'s to be $O(n)$ too, of course.

In summary, my question is: is there a way of defining these probabilities $p_i$ to achieve my goal? Or am I looking for an impossible solution here?

Answer 1

Yes, assuming we can assume that all elements are distinct. Here is one approach. Let $x_1,\dots,x_n$ denote the elements and $w_i$ the weight associated to $x_i$. Consider the following algorithm:

• For $i:=1,2,\dots,n$:
  • Let $p_i=w_i/(w_i+w_{i+1}+\dots+w_n)$.
  • Draw $m_i$ from a Binomial($k$, $p_i$) distribution.
  • Output $m_i$ copies of $x_i$.
  • Set $k := k - m_i$.

This outputs the right distribution.

What is its running time? Well, you can compute the $p_i$'s in linear time, with a single backward scan. Also, you can sample a value $m$ from a Binomial($k$, $p$) distribution in $O(m)$ time. (Draw a uniform random number $q \in [0,1]$, then compute the tail sums of the Binomial, namely $\Pr[m\ge j]$ for $j:=0,1,2,\dots$, until you find the first tail sum that is $\ge q$. It takes a little bit of care to compute the tails sums, but you can do them in $O(1)$ time per tail sum.) Also, we have $m_1+\dots+m_n=k=O(n)$, so it follows that the total running time of this algorithm is $O(n)$.

I don't know whether this will be faster in practice. The constant factors might well dominate the difference between $O(1)$ vs $O(\log n)$.