蓄水池采样算法的python实现_常用算法-蓄水池抽样算法

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()