蓄水池算法

Author Avatar
hatlonely 1月 26, 2018
  • 在其它设备中阅读本文章

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

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

蓄水池算法

问题描述: 给定一个不固定长度的数据集合 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),当所有权重都一致的时候,就和蓄水池算法是一致的了。

这里面有个小问题,就是一开始用来填充样本的数据,其实是等概率的,这样会导致,填充样本的数据权重失效,但是这个问题只在数据集合较小(准确地说 klen(sequence) 比较接近)的情况下才会有比较明显的缺陷,在海量数据集的情况下,这种影响是微乎其微的。

完整代码: https://github.com/hatlonely/algorithm/blob/master/reservoir_sampling.py

参考链接

Reservoir sampling: https://en.wikipedia.org/wiki/Reservoir_sampling

转载请注明出处
本文链接:http://www.hatlonely.com/2018/01/26/蓄水池算法/