常用算法-蓄水池抽样算法

Leetcode上遇到一道题，题目是这样的：

这道题的关键是链表的长度不知道，但是要使随机返回每个元素的概率相等，这一下就难倒我了，如果知道链表的长度k，从0到k中随机选择一个整数就好了呀，可现在不知道长度k的值。不过没关系，蓄水池抽样算法给了我们在不知道长度的情况下，等概率随机抽样的方法。

假设我们要从一大堆元素中选择k个元素，不知道多少，反正很多很多。使用蓄水池抽样算法的具体的思路是：先初始化一个集合，集合中有k个元素，将此集合作为蓄水池。然后从第k+1个元素开始遍历，并且按一定的概率替换掉蓄水池里面的元素。伪代码是这样的：

Init : a reservoir with the size： k  
for i= k+1 to N  
    M=random(1, i);  
    if( M < k)  
     SWAP the Mth value and ith value  
end for

具体描述如下：先将前k个数取出来放入结果集中，然后从第k+1个数开始遍历。假设遍历到第i个数，以k/i的概率替换掉蓄水池中的某个元素即可。

是不是很好理解？接下来，我们来看看，为什么使用这种算法抽取到每个元素的概率是相等的，这里我们使用数学归纳法：

问题描述：

取前k个元素放入蓄水池中。从i=k+1开始，以k／i的概率取第i个元素。若第i个元素被选中，已均等的概率(即1/k)替换蓄水池中的先前被选中的任一元素。

证明：

这样，我们就证明了用蓄水池抽样算法抽取每个元素的概率是相等的。

问题求解

所以，回到上面的leetcode问题，我们得到的答案如下：

# Definition for singly-linked list.
# class ListNode(object):
#     def __init__(self, x):
#         self.val = x
#         self.next = None

class Solution(object):

    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
        """
        result,node,index = self.head,self.head.next,1
        while node:
            if random.randint(0,index) == 0:
                result = node
            node = node.next
            index = index + 1
        return result.val
        


# Your Solution object will be instantiated and called as such:
# obj = Solution(head)
# param_1 = obj.getRandom()