大数据工程师必备之蓄水池抽样算法

做大数据的同学经常会有这样的需求：
给出一个数据流，这个数据流的长度很大或者未知。并且对该数据流中数据只能访问一次。请写出一个随机选择算法，使得数据流中所有数据被选中的概率相等。
或者也可以这么说：
要求从N个元素中随机的抽取k个元素，其中N的大小未知。

很多同学说，擦，这还不简单么，将所有元素保存在一个列表中，然后再随机取k个不就完了么。
好吧，如果你不是专门搞大大大数据的同学，这么说我觉得情有可原。如果你真的是天天跟大大大数据打交道的同学，这么说的话就显得不那么professional了。大数据最重要的特点是什么？一个字，大大大啊！重要的问题必须说三遍。大，那意味着什么？意味着内存放不下呀！所以，我们才需要通过设计精巧的算法，来降低对内存的渴求，达到我们最终的目的。当然可能有土豪会反驳我，俺们公司就是牛逼，俺们就是有钱，内存无限加，128G不够上256G，256G还不够我直接给你加1T内存，反正就是一句话，内存管够！碰到这样的土豪，我也只能淡淡一笑：128G内存到1T内存，即使加上了，也只是翻了10倍。但是现在数据量的增长速度，一般来说可是远远要大于硬件设备的扩容速度的。所以最终，还是需要通过更加精致的算法来解决问题，硬件扩容不是解决问题的根本办法。

前面扯得有点多，下面来看看这个问题怎么解。用到的方法为蓄水池抽样算法（reservoid sampling）。具体的思路是：先初始化一个集合，集合中有k个元素，将此集合作为蓄水池。然后从第k+1个元素开始遍历，并且按一定的概率替换掉蓄水池里面的元素。

来自《The Art of Computer Programming》里的伪代码：

Init : a reservoir with the size： k  
for i= k+1 to N  
    M=random(1, i);  
    if( M < k)  
     SWAP the Mth value and ith value  
end for   

具体描述如下：先将前k个数取出来放入结果集中，然后从第k+1个数开始遍历。假设遍历到第i个数，以 ki 的概率替换掉蓄水池中的某个元素即可。

简单证明一下每个元素出现的概率都是相同的：
假设 n−1 时候成立，即前 n−1 个数据被返回的概率都是 1/n−1 。当前正在读取第 n 个数据，以 1/n 的概率返回它。那么前 n−1 个数据中数据被返回的概率为： (1/(n−1))∗((n−1)/n)=1/n ，假设成立。

参考了hackbuteer1同学的思路，用数学归纳法证明如下。
问题描述：
取前k个元素放入蓄水池中。从 i=k+1 开始，以 ki 的概率取第i个元素。若第i个元素被选中，已均等的概率替换蓄水池中的先前被选中的任一元素。
证明：
当 i=k+1 时，第 k+1 个元素被选中的概率是 kk+1 ，而前k个元素被选中的概率=1 - 被第k+1个元素替换的概率 =1−kk+1×1k=kk+1 ，说明前面 k+1 个元素被取到的概率都是相等的且均为 kk+1 。
假设 i=n 时，前p个元素都以 kn 被选中。
那么当 i=n+1 是，第 n+1 个元素被选中的概率为 kn+1
对于前面的n个元素，每个元素被选中的情况分为两种：1.前面n次已经被选中并且第n+1次时，第n+1个元素没有被选中；2.前面n此已经被选中，第n+1个元素被选中但是没有将其替换掉。不难写出此时的概率为：
kn×(1−kn+1)+kn×(kn+1×(1−1k))=kn+1
由此可见，第n+1步也满足假设条件。所以问题得到证明。

理论说了那么多，一行代码没有，这显然不是我们的风格。一言不合那就直接上代码：

!/usr/bin/env python
#coding:utf-8

import random
import collections

#用蓄水池算法模拟从10个数中随机抽取一个数
def reservoir():
    raw_list = [0,1,2,3,4,5,6,7,8,9]
    ret_num = raw_list[0] #蓄水池初始化。因为只需要抽取一个，所以给一个变量即可

    for i in range(1,10): #从第k+1个元素开始遍历
        m = random.randint(1,i+1) #因为列表下标从0开始，所以随机数上限为i+1而不是i
        if m <= 1:
            ret_num = raw_list[i] #蓄水池里的元素替换

    return ret_num

#抽取十万次，看看最后的结果
def run():
    dic = collections.defaultdict(int)
    for i in range(100000):
        ret_num = reservoir()
        dic[ret_num] += 1

    for k,v in dic.items():
        print k,":",v

run()

将代码run起来以后，看看最后的输出结果：

0 : 9926
1 : 10024
2 : 10056
3 : 10043
4 : 10004
5 : 9826
6 : 10083
7 : 10036
8 : 9843
9 : 10159

从结果可以看出，每个数抽取的次数都在一万次左右，这也就说明，上述代码达到了我们预期的效果！