蓄水池算法

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

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

蓄水池算法

问题描述: 给定一个不固定长度的数据集合 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/蓄水池算法/