最近有个需求，需要从不固定大小的数据集中取固定数量的数据作为样本，有个同学提到了蓄水池算法，于是了解了一下。
蓄水池算法，本身是为了解决海量数据的随机抽样问题，在算法领域应用还是挺广泛的，由于数据本身是有权重，又出现了加权蓄水池算法。
蓄水池算法
问题描述: 给定一个不固定长度的数据集合 sequence，从中等概率地抽取 k 个元素作为样本返回
问题思路: 先把样本填满，然后不断往样本里面等概率替换元素
算法实现
def reservior_sampling(sequence, k):
n = len(sequence)
if k > n:
return sequence
sample = list()
for i in range(k):
sample.append(sequence[i])
for i in range(k, n):
j = random.randint(0, i)
if j >= k:
continue
sample[j] = sequence[i]
return sample
这里需要注意的是往样本里面替换元素的时候，第 i 个元素能被选中用来替换的概率是 k / i + 1，这样就能保证每个元素被选中的机会都是均等的
加权蓄水池算法
问题描述: 给定一个不固定长度的非常大的数据集合 sequence，集合中每个元素包含一个权重 weight，按照权重从集合中抽取 k 个元素返回
问题思路: 和蓄水池算法的思路一样，先把样本填满，然后不断地按照权重替换元素
算法实现
def weighted_reservior_sampling_achao(sequence, k):
n = len(sequence)
if k > n:
return sequence
wsum = 0
sample = list()
for i in range(k):
sample.append(sequence[i])
wsum += sequence[i]['weight'] / k
for i in range(k, n):
wsum += sequence[i]['weight'] / k
p = sequence[i]['weight'] / wsum
j = random.random()
if j <= p:
sample[random.randint(0, k-1)] = sequence[i]
return sample
这里第 i 个元素被选中用来替换的概率是 sequence[i].weight * k / sum(sequence[0:i+1].weight)，当所有权重都一致的时候，就和蓄水池算法是一致的了。
这里面有个小问题，就是一开始用来填充样本的数据，其实是等概率的，这样会导致，填充样本的数据权重失效，但是这个问题只在数据集合较小（准确地说 k 和 len(sequence) 比较接近）的情况下才会有比较明显的缺陷，在海量数据集的情况下，这种影响是微乎其微的。
完整代码: https://github.com/hatlonely/algorithm/blob/master/reservoir_sampling.py
参考链接
Reservoir sampling: https://en.wikipedia.org/wiki/Reservoir_sampling
转载请注明出处
本文链接：http://hatlonely.github.io/2018/01/26/%E8%93%84%E6%B0%B4%E6%B1%A0%E7%AE%97%E6%B3%95/
