今天在網(wǎng)上看題目時(shí)梯轻，發(fā)現(xiàn)一個(gè)十分有趣的算法褪迟，叫蓄水池算法（Reservoir Sampling）彪杉，牽扯到一點(diǎn)概率論問(wèn)題。
題目：給出一個(gè)數(shù)據(jù)流牵咙，這個(gè)數(shù)據(jù)流的長(zhǎng)度很大或者未知。并且對(duì)該數(shù)據(jù)流中數(shù)據(jù)只能訪問(wèn)一次攀唯。請(qǐng)寫(xiě)出一個(gè)隨機(jī)選擇算法洁桌，使得數(shù)據(jù)流中所有數(shù)據(jù)被選中的概率相等。
抽象：從n中取出k個(gè)數(shù)侯嘀，n未知大小另凌，保證最后n中每個(gè)元素被抽取的概率一樣為k/n。
做法：假設(shè)我們從3個(gè)數(shù){1,2,3}中取一個(gè)數(shù)戒幔，那么就要求每個(gè)數(shù)被抽取的概率為1/3吠谢，我們先讀取前2個(gè)數(shù){1,2}，我們以1/2的概率選取其中一個(gè)數(shù)诗茎，加入選擇{1}工坊，接下來(lái)讀取數(shù)字{3}，因?yàn)橐竺總€(gè)數(shù)被選取的概率為1/3敢订，因此我們以1/3的概率選取數(shù)字{3}王污，2/3的概率選取數(shù)字{1}，那么最終數(shù)字1被選取的概率是2/3 * 1/2 = 1/3楚午，同理數(shù)字{2}被選取的概率也是1/3昭齐。
將上述做法n個(gè)數(shù)選擇1個(gè)數(shù)，每次要讀取第n個(gè)數(shù)的時(shí)候矾柜，以1/n的概率保留該數(shù)阱驾，以(n-1)/n的概率保留前面n-1個(gè)數(shù)選取出來(lái)的1個(gè)數(shù)就谜。
再推廣到從n個(gè)數(shù)中選取k個(gè)數(shù)，假設(shè)讀取到第n個(gè)數(shù)時(shí)（n>=k）里覆，以k/n的概率保留概述丧荐，以(n-k)/n的概率保留前n-1個(gè)中選取出來(lái)的k個(gè)數(shù)。
證明：使用上述步驟租谈，從n中讀取k個(gè)數(shù)篮奄，在讀取n個(gè)元素后，n中每一個(gè)被保留下來(lái)的概率都是k/n割去。假設(shè)n = k + i窟却，那么算法證明的就是第i輪選取中，前k+i個(gè)數(shù)每個(gè)數(shù)被保留的概率為k/(k+i)呻逆，其中( 0 <= i <= n - k)夸赫。
用數(shù)學(xué)歸納法來(lái)證明：
- 當(dāng)
i=0是，每個(gè)數(shù)字被選取的概率為k/(k+0)=1咖城，正確茬腿； - 假設(shè)當(dāng)
i-1輪時(shí)，結(jié)論成立宜雀，即前k+i-1中每個(gè)數(shù)被保留的概率為k/(k+i-1)切平； - 第
i輪，因?yàn)槲覀円?code>k/(k+i)的概率選取第k+i個(gè)數(shù)辐董，因此其概率為k/(k+i)悴品，正確；對(duì)于前k+i-1個(gè)數(shù)中的x简烘，其被保留的概率由兩部分組成：
①：第k+i個(gè)數(shù)沒(méi)有被選取到苔严，則x被選取的概率是：i/(k+i) * k/(k+i-1)，其中k/(k+i-1)是2中假設(shè)的條件孤澎，即x在前k+i-1中被保留的概率届氢；
②：第k+i個(gè)數(shù)被選取到，要替換前k+i-1中的數(shù)覆旭，那么x不被替換的概率是：k/(k+i) * k/(k+i-1) * (k-1)/k退子，k/(k+i-1)同①，k-1/k是指型将，在被選取為k個(gè)數(shù)之后絮供，不被替換的概率。
因此茶敏，總概率為(i/(k+i) * k/(k+i-1)) + (k/(k+i) * k/(k+i-1) * (k-1)/k) = (k/(k+i-1)) * (i/(k+i) + (k-1)/(k+i)) = k/k+i壤靶；
得證。
具體算法步驟：
- 從
n個(gè)數(shù)讀取前k個(gè)數(shù)惊搏，保存在集合A中贮乳； - 從第
i個(gè)數(shù)開(kāi)始(k<=i<=n)忧换，每次以k/i的概率選擇是否保留概述，若保留向拆，將隨機(jī)替換k中任一數(shù)亚茬； - 重復(fù)2，直到結(jié)束浓恳。
Leetcode有關(guān)于蓄水池算法的兩道題刹缝，來(lái)看看怎么應(yīng)用：
-
Linked List Random Node
題目要求去鏈表中的任一節(jié)點(diǎn)，使得節(jié)點(diǎn)被選取的概率相等颈将，因此我們利用蓄水池法梢夯，遍歷節(jié)點(diǎn)；
# Definition for singly-linked list.
# class ListNode:
# def __init__(self, x):
# self.val = x
# self.next = None
import random
class Solution:
def __init__(self, head):
"""
@param head The linked list's head.
Note that the head is guaranteed to be not null, so it contains at least one node.
:type head: ListNode
"""
self.head = head
def getRandom(self):
"""
Returns a random node's value.
:rtype: int
"""
node = self.head
select_node = node
i = 1
while node.next:
i += 1
node = node.next
if random.random() <= 1.0/i: # 保證被選取的概率為`1/i`
select_node = node
return select_node.val
-
Random Pick Index
題目要求獲取指定數(shù)字序號(hào)晴圾，每個(gè)數(shù)字在該序列中的序號(hào)被獲取的概率相等颂砸，題目有個(gè)硬性條件，n很大死姚，剛好蓄水池法可以用來(lái)解決人乓。
import random
class Solution:
def __init__(self, nums):
"""
:type nums: List[int]
"""
self.nums = nums
def pick(self, target):
"""
:type target: int
:rtype: int
"""
count = 0
select_index = 0
for index,num in enumerate(self.nums):
if num == target: # 遍歷列表，當(dāng)該值與目標(biāo)值相同時(shí)都毒，才進(jìn)入蓄水池算法
count += 1
if random.random() <= 1.0/count: #滿足概率為1/i
select_index = index
return select_index
