Reservoir sampling is a family of randomized algorithms for choosing a simple random sample, without replacement, of k items from a population of unknown size n in a single pass over the items. The size of the population n is not known to the algorithm and is typically too large for all n items to fit into main memory. The population is revealed to the algorithm over time, and the algorithm cannot look back at previous items. At any point, the current state of the algorithm must permit extraction of a simple random sample without replacement of size k over the part of the population seen so far.

Motivation

Suppose we see a sequence of items, one at a time. We want to keep ten items in memory, and we want them to be selected at random from the sequence. If we know the total number of items n and can access the items arbitrarily, then the solution is easy: select 10 distinct indices i between 1 and n with equal probability, and keep the i-th elements. The problem is that we do not always know the exact n in advance.

Simple: Algorithm R

A simple and popular but slow algorithm, Algorithm R, was created by Jeffrey Vitter.^[1]

Initialize an array ${\displaystyle R}$ indexed from ${\displaystyle 1}$ to ${\displaystyle k}$, containing the first k items of the input ${\displaystyle x_{1},...,x_{k}}$. This is the reservoir.

For each new input ${\displaystyle x_{i}}$, generate a random number j uniformly in ${\displaystyle \{1,...,i\}}$. If ${\displaystyle j\in \{1,...,k\}}$, then set ${\displaystyle R[j]:=x_{i}}$, otherwise, discard ${\displaystyle x_{i}}$.

Return ${\displaystyle R}$ after all inputs are processed.

This algorithm works by induction on ${\displaystyle i\geq k}$.

While conceptually simple and easy to understand, this algorithm needs to generate a random number for each item of the input, including the items that are discarded. The algorithm's asymptotic running time is thus ${\displaystyle O(n)}$. Generating this amount of randomness and the linear run time causes the algorithm to be unnecessarily slow if the input population is large.

This is Algorithm R, implemented as follows:

(* S has items to sample, R will contain the result *)
ReservoirSample(S[1..n], R[1..k])
  // fill the reservoir array
  for i := 1 to k
      R[i] := S[i]
  end

  // replace elements with gradually decreasing probability
  for i := k+1 to n
    (* randomInteger(a, b) generates a uniform integer from the inclusive range {a, ..., b} *)
    j := randomInteger(1, i)
    if j <= k
        R[j] := S[i]
    end
  end
end

Optimal: Algorithm L

If we generate ${\displaystyle n}$ random numbers ${\displaystyle u_{1},...,u_{n}\sim U[0,1]}$ independently, then the indices of the smallest ${\displaystyle k}$ of them is a uniform sample of the k-subsets of ${\displaystyle \{1,...,n\}}$.

The process can be done without knowing ${\displaystyle n}$:

Keep the smallest ${\displaystyle k}$ of ${\displaystyle u_{1},...,u_{i}}$ that has been seen so far, as well as ${\displaystyle w_{i}}$, the index of the largest among them. For each new ${\displaystyle u_{i+1}}$, compare it with ${\displaystyle u_{w_{i}}}$. If ${\displaystyle u_{i+1}<u_{w_{i}}}$, then discard ${\displaystyle u_{w_{i}}}$, store ${\displaystyle u_{i+1}}$, and set ${\displaystyle w_{i+1}}$ to be the index of the largest among them. Otherwise, discard ${\displaystyle u_{i+1}}$, and set ${\displaystyle w_{i+1}=w_{i}}$.

Now couple this with the stream of inputs ${\displaystyle x_{1},...,x_{n}}$. Every time some ${\displaystyle u_{i}}$ is accepted, store the corresponding ${\displaystyle x_{i}}$. Every time some ${\displaystyle u_{i}}$ is discarded, discard the corresponding ${\displaystyle x_{i}}$.

This algorithm still needs ${\displaystyle O(n)}$ random numbers, thus taking ${\displaystyle O(n)}$ time. But it can be simplified.

First simplification: it is unnecessary to test new ${\displaystyle u_{i+1},u_{i+2},...}$ one by one, since the probability that the next acceptance happens at ${\displaystyle u_{i+l}}$ is ${\displaystyle (1-u_{w_{i}})^{l-1}-(1-u_{w_{i}})^{l}}$, that is, the interval ${\displaystyle l}$ of acceptance follows a geometric distribution.

Second simplification: it's unnecessary to remember the entire array of the smallest ${\displaystyle k}$ of ${\displaystyle u_{1},...,u_{i}}$ that has been seen so far, but merely ${\displaystyle w}$, the largest among them. This is based on three observations:

• Every time some new ${\displaystyle x_{i+1}}$ is selected to be entered into storage, a uniformly random entry in storage is discarded.
• ${\displaystyle u_{i+1}\sim U[0,w]}$
• ${\displaystyle w_{i+1}}$ has the same distribution as ${\displaystyle \max\{u_{1},...,u_{k}\}}$, where all ${\displaystyle u_{1},...,u_{k}\sim U[0,w]}$ independently. This can be sampled by first sampling ${\displaystyle u\sim U[0,1]}$, then taking ${\displaystyle w\cdot u^{\frac {1}{k}}}$.

This is Algorithm L,^[2] which is implemented as follows:

(* S has items to sample, R will contain the result *)
ReservoirSample(S[1..n], R[1..k])
  // fill the reservoir array
  for i = 1 to k
      R[i] := S[i]
  end

  (* random() generates a uniform (0,1) random number *)
  W := exp(log(random())/k)

  while i <= n
      i := i + floor(log(random())/log(1-W)) + 1
      if i <= n
          (* replace a random item of the reservoir with item i *)
          R[randomInteger(1,k)] := S[i]  // random index between 1 and k, inclusive
          W := W * exp(log(random())/k)
      end
  end
end

This algorithm computes three random numbers for each item that becomes part of the reservoir, and does not spend any time on items that do not. Its expected running time is thus ${\displaystyle O(k(1+\log(n/k)))}$,^[2] which is optimal.^[1] At the same time, it is simple to implement efficiently and does not depend on random deviates from exotic or hard-to-compute distributions.

With random sort

If we associate with each item of the input a uniformly generated random number, the k items with the largest (or, equivalently, smallest) associated values form a simple random sample.^[3] A simple reservoir-sampling thus maintains the k items with the currently largest associated values in a priority queue.

(*
  S is a stream of items to sample
  S.Current returns current item in stream
  S.Next advances stream to next position
  min-priority-queue supports:
    Count -> number of items in priority queue
    Minimum -> returns minimum key value of all items
    Extract-Min() -> Remove the item with minimum key
    Insert(key, Item) -> Adds item with specified key
 *)
ReservoirSample(S[1..?])
  H := new min-priority-queue
  while S has data
    r := random()   // uniformly random between 0 and 1, exclusive
    if H.Count < k
      H.Insert(r, S.Current)
    else
      // keep k items with largest associated keys
      if r > H.Minimum
        H.Extract-Min()
        H.Insert(r, S.Current)
    end
    S.Next
  end
  return items in H
end

The expected running time of this algorithm is ${\displaystyle O(n+k\log k\log(n/k))}$ and it is relevant mainly because it can easily be extended to items with weights.

Weighted random sampling

This method, also called sequential sampling, is incorrect in the sense that it does not allow to obtain a priori fixed inclusion probabilities. Some applications require items' sampling probabilities to be according to weights associated with each item. For example, it might be required to sample queries in a search engine with weight as number of times they were performed so that the sample can be analyzed for overall impact on user experience. Let the weight of item i be ${\displaystyle w_{i}}$, and the sum of all weights be W. There are two ways to interpret weights assigned to each item in the set:^[4]

1. In each round, the probability of every unselected item to be selected in that round is proportional to its weight relative to the weights of all unselected items. If X is the current sample, then the probability of an item ${\displaystyle i\notin X}$ to be selected in the current round is ${\displaystyle \textstyle w_{i}/(W-\sum _{j\in X}{w_{j}})}$.
2. The probability of each item to be included in the random sample is proportional to its relative weight, i.e., ${\displaystyle w_{i}/W}$. Note that this interpretation might not be achievable in some cases, e.g., ${\displaystyle k=n}$.

(*
  S is a stream of items to sample
  S.Current returns current item in stream
  S.Weight  returns weight of current item in stream
  S.Next advances stream to next position
  The power operator is represented by ^
  min-priority-queue supports:
    Count -> number of items in priority queue
    Minimum() -> returns minimum key value of all items
    Extract-Min() -> Remove the item with minimum key
    Insert(key, Item) -> Adds item with specified key
 *)
ReservoirSample(S[1..?])
  H := new min-priority-queue
  while S has data
    r := random() ^ (1/S.Weight)   // random() produces a uniformly random number in (0,1)
    if H.Count < k
      H.Insert(r, S.Current)
    else
      // keep k items with largest associated keys
      if r > H.Minimum
        H.Extract-Min()
        H.Insert(r, S.Current)
      end
    end
    S.Next
  end
  return items in H
end

(*
  S is a stream of items to sample
  S.Current returns current item in stream
  S.Weight  returns weight of current item in stream
  S.Next advances stream to next position
  The power operator is represented by ^
  min-priority-queue supports:
    Count -> number of items in the priority queue
    Minimum -> minimum key of any item in the priority queue
    Extract-Min() -> Remove the item with minimum key
    Insert(Key, Item) -> Adds item with specified key
 *)
ReservoirSampleWithJumps(S[1..?])
  H := new min-priority-queue
  while S has data and H.Count < k
    r := random() ^ (1/S.Weight)   // random() produces a uniformly random number in (0,1)
    H.Insert(r, S.Current)
    S.Next
  end
  X := log(random()) / log(H.Minimum) // this is the amount of weight that needs to be jumped over
  while S has data
    X := X - S.Weight
    if X <= 0
      t := H.Minimum ^ S.Weight
      r := random(t, 1) ^ (1/S.Weight) // random(x, y) produces a uniformly random number in (x, y)
    
      H.Extract-Min()
      H.Insert(r, S.Current)

      X := log(random()) / log(H.Minimum)
    end
    S.Next
  end
  return items in H
end

(*
  S has items to sample, R will contain the result
  S[i].Weight contains weight for each item
 *)
WeightedReservoir-Chao(S[1..n], R[1..k])
  WSum := 0
  // fill the reservoir array
  for i := 1 to k
      R[i] := S[i]
      WSum := WSum + S[i].Weight
  end
  for i := k+1 to n
    WSum := WSum + S[i].Weight
    p := S[i].Weight / WSum // probability for this item
    j := random();          // uniformly random between 0 and 1
    if j <= p               // select item according to probability
        R[randomInteger(1,k)] := S[i]  //uniform selection in reservoir for replacement
    end
  end
end

(*
  S has items to sample, R will contain the result
  S[i].Weight contains weight for each item
 *)
WeightedReservoir-Chao(S[1..n], R[1..k])
  WSum := 0
  // fill the reservoir array
  for i := 1 to k
      R[i] := S[i]
      WSum := WSum + S[i].Weight
  end
  j := random() // uniformly random between 0 and 1
  pNone := 1 // probability that no item has been selected so far (in this jump)
  for i := k+1 to n
    WSum := WSum + S[i].Weight
    p := S[i].Weight / WSum // probability for this item
    j -= p * pNone
    pNone := pNone * (1 - p)
    if j <= 0
        R[randomInteger(1,k)] := S[i]  //uniform selection in reservoir for replacement
        j = random()
        pNone := 1
    end
  end
end

Relation to Fisher–Yates shuffle

Suppose one wanted to draw k random cards from a deck of cards. A natural approach would be to shuffle the deck and then take the top k cards. In the general case, the shuffle also needs to work even if the number of cards in the deck is not known in advance, a condition which is satisfied by the inside-out version of the Fisher–Yates shuffle:^[9]

(* S has the input, R will contain the output permutation *)
Shuffle(S[1..n], R[1..n])
  R[1] := S[1]
  for i from 2 to n do
      j := randomInteger(1, i)  // inclusive range
      R[i] := R[j]
      R[j] := S[i]
  end
end

Note that although the rest of the cards are shuffled, only the first k are important in the present context. Therefore, the array R need only track the cards in the first k positions while performing the shuffle, reducing the amount of memory needed. Truncating R to length k, the algorithm is modified accordingly:

(* S has items to sample, R will contain the result *)
ReservoirSample(S[1..n], R[1..k])
  R[1] := S[1]
  for i from 2 to k do
      j := randomInteger(1, i)  // inclusive range
      R[i] := R[j]
      R[j] := S[i]
  end
  for i from k + 1 to n do
      j := randomInteger(1, i)  // inclusive range
      if (j <= k)
          R[j] := S[i]
      end
  end
end

Since the order of the first k cards is immaterial, the first loop can be removed and R can be initialized to be the first k items of the input. This yields Algorithm R.

Limitations

Reservoir sampling makes the assumption that the desired sample fits into main memory, often implying that k is a constant independent of n. In applications where we would like to select a large subset of the input list (say a third, i.e. ${\displaystyle k=n/3}$), other methods need to be adopted. Distributed implementations for this problem have been proposed.^[10]

References