蓄水池算法

今天在网上看题目时，发现一个十分有趣的算法，叫蓄水池算法（Reservoir Sampling），牵扯到一点概率论问题。

题目：给出一个数据流，这个数据流的长度很大或者未知。并且对该数据流中数据只能访问一次。请写出一个随机选择算法，使得数据流中所有数据被选中的概率相等。

抽象：从n中取出k个数，n未知大小，保证最后n中每个元素被抽取的概率一样为k/n。

做法：假设我们从3个数{1,2,3}中取一个数，那么就要求每个数被抽取的概率为1/3，我们先读取前2个数{1,2}，我们以1/2的概率选取其中一个数，加入选择{1}，接下来读取数字{3}，因为要求每个数被选取的概率为1/3，因此我们以1/3的概率选取数字{3}，2/3的概率选取数字{1}，那么最终数字1被选取的概率是2/3 * 1/2 = 1/3，同理数字{2}被选取的概率也是1/3。

将上述做法n个数选择1个数，每次要读取第n个数的时候，以1/n的概率保留该数，以(n-1)/n的概率保留前面n-1个数选取出来的1个数。

再推广到从n个数中选取k个数，假设读取到第n个数时（n>=k），以k/n的概率保留概述，以(n-k)/n的概率保留前n-1个中选取出来的k个数。

证明：使用上述步骤，从n中读取k个数，在读取n个元素后，n中每一个被保留下来的概率都是k/n。假设n = k + i，那么算法证明的就是第i轮选取中，前k+i个数每个数被保留的概率为k/(k+i)，其中( 0 <= i <= n - k)。
用数学归纳法来证明：

1. 当i=0是，每个数字被选取的概率为k/(k+0)=1，正确；
2. 假设当i-1轮时，结论成立，即前k+i-1中每个数被保留的概率为k/(k+i-1)；
3. 第i轮，因为我们以k/(k+i)的概率选取第k+i个数，因此其概率为k/(k+i)，正确；对于前k+i-1个数中的x，其被保留的概率由两部分组成：
   ①：第k+i个数没有被选取到，则x被选取的概率是：i/(k+i) * k/(k+i-1)，其中k/(k+i-1)是2中假设的条件，即x在前k+i-1中被保留的概率；
   ②：第k+i个数被选取到，要替换前k+i-1中的数，那么x不被替换的概率是：k/(k+i) * k/(k+i-1) * (k-1)/k，k/(k+i-1)同①，k-1/k是指，在被选取为k个数之后，不被替换的概率。
   因此，总概率为(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；

得证。

具体算法步骤：

1. 从n个数读取前k个数，保存在集合A中；
2. 从第i个数开始(k<=i<=n)，每次以k/i的概率选择是否保留概述，若保留，将随机替换k中任一数；
3. 重复2，直到结束。

Leetcode有关于蓄水池算法的两道题，来看看怎么应用：

1. Linked List Random Node
   题目要求去链表中的任一节点，使得节点被选取的概率相等，因此我们利用蓄水池法，遍历节点；

# 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

2. Random Pick Index
   题目要求获取指定数字序号，每个数字在该序列中的序号被获取的概率相等，题目有个硬性条件，n很大，刚好蓄水池法可以用来解决。

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:      # 遍历列表，当该值与目标值相同时，才进入蓄水池算法
                count += 1
                if random.random() <= 1.0/count:      #满足概率为1/i
                    select_index = index

        return select_index