Reservoir sampling
==================

Reservoir sampling is an algorithm to choose a random :math:`k`-size
subset of :math:`N` elements, where :math:`N` is very large and
possibly unknown. Throughout we assume that :math:`k \le N`.


Warmup: small :math:`N`
-----------------------

Suppose we have a list of :math:`N` elements from which we would like
a random sample of size :math:`k` drawn uniformly at random. If
:math:`N` is known and small enough that the list of :math:`N`
elements can fit into main memory, we could find our subset as
follows. Generate a random permutation of the :math:`N` elements using
a `Fisher-Yates`_ shuffle, and then choose the first :math:`k`
elements from the list. Since we began with a random permutation of
all :math:`N` elements, the probability of any :math:`k`-size subset
being in the first :math:`k` elements is the same. The Fisher-Yates
shuffle is shown below.

.. _Fisher-Yates: https://en.wikipedia.org/wiki/Fisher%E2%80%93Yates_shuffle

.. code-block:: none

   SHUFFLE(xs)
     let n = length xs
     for i = 1 to n - 1
       let j = draw uniformly at random from {i, i + 1, ..., n}
       swap xs[i], xs[j]

This algorithm takes :math:`O(n)` time to shuffle :math:`n`
elements. We can improve the runtime to :math:`O(k)` where we only
need to shuffle the first :math:`k` elements as follows.

.. code-block:: none

   SHUFFLE(xs, k)
     let n = length xs
     for i = 1 to k
       let j = draw uniformly at random from {i, i + 1, ..., n}
       swap xs[i], xs[j]

The sampling algorithm proposed so far takes :math:`O(k)` time and
:math:`O(N)` space, and requires knowledge of the value of :math:`N`.


Large or unknown :math:`N`
--------------------------

Our algorithm proceeds as follows. As we process each element of the
list we maintain a reservoir of :math:`k` elements which represents
our current choice of a random size :math:`k` subset should the list
have no additional elements. Whenever we encounter a new element in
the list, we replace an element of the reservoir with the new element
with the appropriate probability. The pseudocode is as follows.

.. code-block:: none

   SAMPLE(k, list)
     // Initialization; get the first k elements
     for i = 1 to k
       let r[i] = list[i]

     // Inductive step
     let i = k + 1
     while list has an element at i
       let x = list[i]
       with probability k / i
         let j = draw uniformly at random from {1, 2, ..., k}
         let r[j] = x
       let i = i + 1

To prove that the algorithm is correct, we maintain the following
inductive invariant: after viewing :math:`i` elements, our reservoir
contains a size :math:`k` subset of the first :math:`i` elements drawn
uniformly at random.

To begin, we note that the invariant holds for :math:`i = k` because
of our initialization routine. We now show that it holds for general
:math:`i \ge k`. Let :math:`i > k` be given, and assume the invariant
has held up to processing element :math:`i - ``. Since there are
:math:`i` elements in total and we are selecting :math:`k`, the
probability that any one element is in the subset is :math:`k /
i`. Our algorithm clearly chooses the new element for inclusion with
the appropriate probability. Additionally, each of the first
:math:`i - 1` elements are in the subset with probability :math:`k /
(i - 1)` by our invariant. Any such element is replaced by the new
element with probability :math:`(k / i) / k = 1 / i`. Thus the
probability any one of the first :math:`i - 1` elements are in the
subset after this iteration is :math:`k / (i - 1) \cdot (1 - 1 / i) =
k / i`, as required.

We have shown that the marginal probability of any one element being
in the subset is correct; however, we need to show that the
probability of any :math:`k`-size subset being chosen is equal. Let
any size :math:`k` subset of the first :math:`i` elements be
given. We need that the probability of choosing this subset is
:math:`\binom{i}{k}^{-1}`. There are two cases to consider: either
element :math:`i` is in the subset, or not. In the first case we have
from our invariant that the probability of this subset being chosen is

.. math::

   \binom{i - 1}{k}^{-1} (1 - k / i) = \frac{k! (i - k - 1)!}{(i -
   1)!} \cdot \frac{i - k}{i} = \frac{k! (i - k)!}{i!} =
   \binom{i}{k}^{-1},

as required. Otherwise, we need all but the :math:`i\text{th}` element
to be in the previous subset, which happens with probability
:math:`(i - k) \binom{i - 1}{k}^{-1}` since there are :math:`i - k`
:math:`k`-size subsets containing the first :math:`k - 1` elements of
our group. We then need to replace the one bad element with the new
element, which happens with probability :math:`(k / i) \cdot (1 /
k)`. In total, the probability we have the given subset after the
:math:`i\text{th}` element has been processed is

.. math::

   (i - k) \binom{i - 1}{k}^{-1} \frac{k}{i} \cdot \frac{1}{k} =
   \frac{(i - k)}{i} \cdot \frac{k! (i - 1 - k)!}{(i - 1)!} = \frac{k!
   (i - k)!}{i!} = \binom{i}{k}^{-1},

as required.