蓄水池算法

Posted 2020-10-22 hatlonely

最近有个需求，需要从不固定大小的数据集中取固定数量的数据作为样本，有个同学提到了蓄水池算法，于是了解了一下。

蓄水池算法，本身是为了解决海量数据的随机抽样问题，在算法领域应用还是挺广泛的，由于数据本身是有权重，又出现了加权蓄水池算法。

蓄水池算法

问题描述: 给定一个不固定长度的数据集合 sequence，从中等概率地抽取 k 个元素作为样本返回

问题思路: 先把样本填满，然后不断往样本里面等概率替换元素

算法实现

1. def reservior_sampling(sequence, k):
2. n = len(sequence)
3. if k > n:
4. return sequence
6. sample = list()
7. for i in range(k):
8. sample.append(sequence[i])
10. for i in range(k, n):
11. j = random.randint(0, i)
12. if j >= k:
13. continue
14. sample[j] = sequence[i]
16. return sample

这里需要注意的是往样本里面替换元素的时候，第 i 个元素能被选中用来替换的概率是 k / i + 1，这样就能保证每个元素被选中的机会都是均等的

加权蓄水池算法

问题描述: 给定一个不固定长度的非常大的数据集合 sequence，集合中每个元素包含一个权重 weight，按照权重从集合中抽取 k 个元素返回

问题思路: 和蓄水池算法的思路一样，先把样本填满，然后不断地按照权重替换元素

算法实现

1. def weighted_reservior_sampling_achao(sequence, k):
2. n = len(sequence)
3. if k > n:
4. return sequence
6. wsum = 0
7. sample = list()
8. for i in range(k):
9. sample.append(sequence[i])
10. wsum += sequence[i][‘weight‘] / k
12. for i in range(k, n):
13. wsum += sequence[i][‘weight‘] / k
14. p = sequence[i][‘weight‘] / wsum
15. j = random.random()
16. if j <= p:
17. sample[random.randint(0, k-1)] = sequence[i]
19. 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/蓄水池算法/