水塘抽样的原理及python实现

同步更新于知乎：https://zhuanlan.zhihu.com/p/107889958

最近在做整理，偶尔看到公众号的一篇关于随机抽样的分享，这个算法面试中经常会问到，特此总结一下。

网上关于这块解释并不清晰，主要参考如下，个人感觉写的比他更通俗易懂，哈哈。

https://zhuanlan.zhihu.com/p/107793995

解决什么问题

主要用于解决大数据流中的随机抽样问题，即：当内存有限，数据长度很大，甚至未知，那么如何从中随机选取k个数据，并且要求是等概率

算法核心

水塘抽样的核心是，只遍历一次，每次都考虑一个问题：当前元素是否被选中，选中后替换之前选中的哪一个元素。

知道这个思想，只需要记住结论就可以了：

如果要随机选择K个元素，那么在遍历到第i个元素时，以k/i的概率选择该元素。

代码实现

转换成代码时，只需要考虑如何构造k/i的概率即可，详细见下面python代码。

两种情况：

最简单的情况，当K = 1时：

def reservoirSampling1(arr):
    # 初始化把第1个元素放进结果中
    res = arr[0]
    i = 1  # 从第2个元素开始考虑
    while i < len(arr):
        # 概率1/i选取第i个元素作为结果
        r = random.randint(0, i)
        if r == 0:
            res = arr[i]
        i += 1
    # 这样，遍历一遍就可以等概率的选择一个数输出
    return res

K > 1 时：

def reservoirSamplingk(arr, k):
    res = arr[:k]
    i = k
    while i < len(arr):
        # 以k/i的概率选取第i个元素，用来等概率的替换之前选中的1个元素
        r = random.randint(0, i)
        if r < k:  # 小于k的概率就是k/i，替换res中第r个已选中的数
            res[r] = arr[i]
        i += 1
    return res

数学原理

知道结论后，如果想更进一步，可以有简单的数学推导：

在遍历到第i个元素时，如果以k/i的概率保留，并且使得最终被选中，那么之后的n-i个元素都要不被替换才行。

注意，在第i+1个元素时，第i个元素不被替换的概率，也就是第i+1个元素不被保留的概率是1-k/(i+1)*1/k，后面这个1/k表示在result数组中等概率被替换到。

即：
$$\frac{k}{i}\ast (1-\frac{k}{i+1}\ast \frac{1}{k})\ast (1-\frac{k}{i+2}\ast \frac{1}{k})\ast ...\ast (1-\frac{k}{n}\ast \frac{1}{k})=k/n$$
很明显，每个元素被都是以概率k/n，等概率的被选择到。

由于水塘抽样必须遍历一次所有数据，所以时间复杂度是O(n)。