最近有個(gè)需求，需要從不固定大小的數(shù)據(jù)集中取固定數(shù)量的數(shù)據(jù)作為樣本，有個(gè)同學(xué)提到了蓄水池算法，于是了解了一下。

蓄水池算法，本身是為了解決海量數(shù)據(jù)的隨機(jī)抽樣問(wèn)題，在算法領(lǐng)域應(yīng)用還是挺廣泛的，由于數(shù)據(jù)本身是有權(quán)重，又出現(xiàn)了加權(quán)蓄水池算法。

蓄水池算法

問(wèn)題描述: 給定一個(gè)不固定長(zhǎng)度的數(shù)據(jù)集合 sequence，從中等概率地抽取 k 個(gè)元素作為樣本返回

問(wèn)題思路: 先把樣本填滿，然后不斷往樣本里面等概率替換元素

算法實(shí)現(xiàn)

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

這里需要注意的是往樣本里面替換元素的時(shí)候，第 i 個(gè)元素能被選中用來(lái)替換的概率是 k / i + 1，這樣就能保證每個(gè)元素被選中的機(jī)會(huì)都是均等的

加權(quán)蓄水池算法

問(wèn)題描述: 給定一個(gè)不固定長(zhǎng)度的非常大的數(shù)據(jù)集合 sequence，集合中每個(gè)元素包含一個(gè)權(quán)重 weight，按照權(quán)重從集合中抽取 k 個(gè)元素返回

問(wèn)題思路: 和蓄水池算法的思路一樣，先把樣本填滿，然后不斷地按照權(quán)重替換元素

算法實(shí)現(xiàn)

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 個(gè)元素被選中用來(lái)替換的概率是 sequence[i].weight * k / sum(sequence[0:i+1].weight)，當(dāng)所有權(quán)重都一致的時(shí)候，就和蓄水池算法是一致的了。

這里面有個(gè)小問(wèn)題，就是一開始用來(lái)填充樣本的數(shù)據(jù)，其實(shí)是等概率的，這樣會(huì)導(dǎo)致，填充樣本的數(shù)據(jù)權(quán)重失效，但是這個(gè)問(wèn)題只在數(shù)據(jù)集合較?。?zhǔn)確地說(shuō) k 和 len(sequence) 比較接近）的情況下才會(huì)有比較明顯的缺陷，在海量數(shù)據(jù)集的情況下，這種影響是微乎其微的。

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

參考鏈接

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

轉(zhuǎn)載請(qǐng)注明出處
本文鏈接：http://hatlonely.github.io/2018/01/26/%E8%93%84%E6%B0%B4%E6%B1%A0%E7%AE%97%E6%B3%95/