蓄水池抽样算法及其python实现

蓄水池抽样算法

蓄水池抽样是个很有趣的问题，这个问题的来源是关于等概率抽样的一种思考，问题是，如何能在不知道总体对象数量（或者数量巨大）的情况下抽取k个对象，使得每个对象被抽取到的概率相同。解决问题的思路：考虑最终一定要抽取到k个对象，所以先任意抽出k个，然后对剩下的对象分别以某种概率概率，使得最终每个对象被抽到的概率相同。（根据分步分类计数原理即可证明）

算法流程：

1. 输入：长度为N的数组L（N未知或者很大）；输出：被等可能抽出的长度为k的数组l
2. 对输入L取前k个数组成的数组作为蓄水池
3. 对于L的第i(i=k+1,k+2,…,N)个数，任取r为0~i-1之间的整数，若r>k-1,则不进行替换，若r<=k-1,则用第i个数去替换蓄水池中第r个数
4. 遍历一遍L，取到的l中的每个元素都是以概率k/n等可能取到的

python实现如下：

#coding=utf-8'''Created on 2016年9月15日Reservoir sampling的python实现@author: whz'''import numpy as npimport matplotlib.pyplot as pltdef pool(L,k):    arr = L[:k]    for i,e in enumerate(L[k:]):        r = np.random.randint(0,k+i+1)        if r<=k-1:            arr[r] = e    return arrdef count(q,n):    L = array("d")    l1 = array("d")    l2 = array("d")    for e in q:        L.append(e)    for e in range(n):        l1.append(L.count(e))    for e in l1:        l2.append(e/sum(l1))    return l1,l2if __name__ == '__main__':    a = np.array([[i]*10000000 for i in range(3)])#生成等量的0，1，2    a.shape = 1,-1    L = a[0]    k = 300000#设置要抽取的样本的数量，一般远小于总体数量    p = pool(L, k)    l1 = ['value=%d'% x for x in range(3)]    plt.pie(count(p,3)[0],labels=l1,labeldistance=0.1,autopct='%1.2f%%')    plt.title("Reservoir sampling")    plt.show()

以上代码设计了一个实验来验证蓄水池算法的准确性：总体是由0,1,2各一千万个组成的三千万个数的数组，从中抽取三十万个数，来观察抽出样本的统计特征，如下图：

验证了这种抽样方法在一定精度上能够实现等概率抽样。