【演算法】蓄水池抽樣
阿新 • • 發佈:2019-01-04
例題:有一個機器按自然數序列的方式吐出球,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