1. 程式人生 > >【演算法】蓄水池抽樣

【演算法】蓄水池抽樣

例題:有一個機器按自然數序列的方式吐出球,1號球,2號球......。現有一個袋子,袋子裡最多隻能裝下k個球,並且除袋子以外沒有更多的空間,球扔掉不能放回。設計一種選擇方式,使得當機器吐出第N號球時,袋子中的球數是k個,同時可以保證從1號球到N號球中的每一個被選中進袋子的概率都是k/N。

具體例子:

有一個只能裝下10個球的袋子:

當吐出100個球時,袋子裡有10個球,且1~100號球中每個被選中放入袋子的概率都是10/100;

當吐出1000個球時,袋子裡有10個球,且1~1000號球中每個被選中放入袋子的概率都是10/1000;

當吐出N個球時(N>10),袋子裡有10個球,且1~N號球中每個被選中放入袋子的概率都是10/N;

蓄水池演算法步驟:

1. 處理1~k號球時,直接放入袋子;

2. 處理第i號球(i > k)時,以 k/i 概率決定是否將第i號球放入袋子。若不放入,直接扔掉。若放入,則從袋子裡隨機選一個球扔掉,然後把第i號球放進袋子。

 

import random;

class Bag:
    selected = []
    def __init__(self):
        self.selected = []
    # 每次拿一個球都會呼叫這個函式,N表示第i次呼叫
    def carryBalls(self, N, k):
        if N <= k:
            return self.selected.append(N)
        if random.randint(1,N) <= k:
            index = random.randint(0,k-1)
            self.selected[index] = N
        
        return self.selected