蓄水池抽样算法

引子

• 问题：流式数据（Streaming Data），数据长度为n但不知道，如何从中等概率随机选择一个数据？
  
• 解法：我们以概率1选择第一个数据，以1/2的概率选择第二个数据，以此类推，以1/m的概率选择第m个对象（如果后面某一数据一旦选中，替换掉以前选中的数据）。当所有数据流过时，每个对象具有相同的被选中的概率1/n。
  
• 证明：$P_m = \frac{1}{m} \cdot \frac{m}{m+1} \cdot \frac{m+1}{m+2} \cdot *** \cdot \frac{n-2}{n-1} \cdot \frac{n-1}{n} = \frac{1}{n}$

• Python代码

  import random
  def sample(n):
      res = 1
      for i in range(1, n+1):
          num = random.randint(1, i):
          if i == num:
              res = i
      return res

蓄水池抽样

• 问题：流式数据（Streaming Data），数据长度为n但不知道，如何从中等概率随机选择k（k < n）个数据？
• 解法：先选中前k个数据（当作一个蓄水池），然后以k/(k+1)的概率选择第k+1个数据，以k/(k+2)的概率选择第k+2个数据，以此类推以k/m的概率选择第m个数据，一旦后面有某个数据被选中，则随机替换蓄水池中的一个数据。最终每个数据被选中的概率均为k/n。
• 证明： $$\begin{align} P_m &= \frac{k}{m} (\frac{m+1-k}{m+1} +\frac{k}{m+1} \frac{k-1}{k}) (\frac{m+2-k}{m+2} +\frac{k}{m+2}\frac{k-1}{k}) ...(\frac{n-k}{n} +\frac{k}{n}\frac{k-1}{k}) \\ &= \frac{k}{n} \end{align}$$
• Python代码

import random
def sample(n, k):
    res = [0] * k
    for i range(k):
        res[i] = i+1
    for i range(k, n):
        num = random.randint(0, i)
        if num < k:
            res[num]  = i+1
    return res

参考

• 蓄水池抽样
• 蓄水池抽样算法(Reservoir Sampling)