蓄水池抽样算法是比较经典的算法，可以高效的在大数据中等概率且不重复采样，比如从一个大文件中等概率抽取m行内容打印出来。
说白了就是需要 从n个数中等概率且不重复取出m个数 (其中m<=n)
先贴出最终代码如下(数组下标从1开始可以更好阐述问题，所以我使用lua代码表示)
--1.假设dataStream,是一个长度为n的数组
--2.定义pool数组用于存放最终采样出来的m个数
local pool = {}
--3.算法
for i = 1,n do
if i <= m then
pool[i] = dataStream[i]
else
local r1 = random(0,1)
if r1 <= m/i then --这里是 m/i 的概率
local r2 = random(1,m) --这里是 1/m 的概率
pool[r2] = dataStream[i]
end
end
end
一、验证分析
我们先假设 m = 1
当i = 1时 1<=m, 所以dataStream[1]直接被存到pool 中，它的概率是1
当i = 2时 2>m, 所以开始随机0~1之间的一个数，如果结果r1 满足 1/2 的概率，那最终它会被写入pool，
这说明了两件事
1.dataStream[2],被写入的概率为 1/2
2.此时dataStream[1]的概率是多少？
上一次(i=1时)的dataStream[1]的写入pool的概率是1，
这一次(i=2时)dataStream[1]不被替换的概率是（1-(1/2)x1）
这两步同时满足时dataStream[1]才被最终保留(真正写入pool)，它的概率是 1x(1-(1/2)x1) = 1/2，发现与dataStream[2]的概率相同
当i = 3时 3>m,所以开始随机0~1之间的一个数，如果结果r1 满足 1/3 的概率，那最终它会被写入pool，
这说明了三件事
1.dataStream[3],被写入的概率为 1/3
2.此时dataStream[1]的概率是多少？
上一次(i=1时)的dataStream[1]的写入pool的概率是1/2
这一次(i=3时)dataStream[1]不被替换的概率是（1-(1/3)x1）
这两步同时满足时dataStream[1]才被最终保留(真正写入pool)，它的概率是 1/2x(1-(1/3)x1) = 1/3，发现与dataStream[3]的概率相同
3.此时dataStream[2]的概率是多少？
它的概率计算与dataStream[1]一样，所以也是1/3
根据数学归纳法，可得当i=n时(遍历到最后)，使用这个算法可以使每一项被留在pool中的概率都为1/n (pool的容量m=1)
继续往下推
当m=m时(m为<n的正整数)
dataStream[i],被写入的概率为 m/i
上一次(i=i-1时)的dataStream[i]的写入pool的概率是m/(i-1)
这一次(i=i 时)dataStream[1]不被替换的概率是 1 - (m/i)x(1/m)
这两步同时满足时dataStream[1]才被最终保留(真正写入pool)它的概率是 m/(i-1) x (1 - (m/i)x(1/m)) = m/i，发现与dataStream[i]的概率相同
dataStream[2]到dataStream[n-1]的概率计算与dataStream[1]是相同的，结果都是m/i
最终结论：当i = n时，使用这个算法可以使每一项被留在pool中的概率都为m/n
二、代码优化
上面的代码使用了两个随机函数(目的是为了更好的理解概率计算)，事实上我们也可以通过一个随机来达到同样的效果，具体代码如下
--1.假设dataStream,是一个长度为n的数组
--2.定义pool数组用于存放最终采样出来的m个数
local pool = {}
--3.算法
for i = 1,n do
if i <= m then
pool[i] = dataStream[i]
else
local r = math.random(1,i)
if r <= m then
pool[r] = dataStream[i]
end
end
end
三、最后我们来看一下这个算法的时间复杂度与空间复杂度
时间复杂度O(n) : 一次遍历搞定问题
空间复杂度 O(m) : 也就是pool的大小
文章评论