1. 问题
给定一个单链表,随机返回一个结点,要求每个结点被选中的概率相等。
2. 思路
在一个给定长度的数组中等概率抽取一个数,可以简单用随机函数random.randint(0, n-1)得到索引来抽取。
本题是给定了链表,当然也好做,可以事先遍历一次求长度,每次要取的时候随机求索引,然后遍历一次。
时间复杂度O(n),空间复杂度O(1)
或者事先把数据放到数组中,每次要取的时候随机求索引,然后直接取到对应的数。
时间复杂度O(1),空间复杂度O(n)
(延伸一下)如果对于长度未知,会不断增加的数据流呢?可以使用蓄水池采样(Reservoir Sampling)的方法。如果我们要从n个数(这个n会不断增加)中等概率地抽取k个数,做法如下:
(1)先取数据流的前k个数,保存在数组reservoir中。
(2)对于第i个数(k+1 <= i <= n),以k/i的概率选择是否保留第i个数,如果第i个数被选中,则从reservoir中随机选择一个数,用第i个数代替它。
(3)重复迭代第二步,reservoir中的k个数就是我们要的结果。
蓄水池采样的证明
为什么这么做可以保证等概率抽取,这里证明一下。我们现在要做的事是等概率的抽取k个数。
(1)当只有k个数的时, 每个数被抽取的概率是k/k(也就是1啦),k个数都放到蓄水池中。
(2)假设现在增加了一个数,第k+1个数,用k/(k+1)的概率选择是否保留,此时第k+1个数被保留的概率是k/(k+1)。
对于蓄水池中的k个数,它们被留下的概率是多少呢?它们原来的概率都是1,但是现在新来了一个数据,每个数据都面临被淘汰的风险。
淘汰的概率为(第k+1个数被选中的概率)乘以(每个数据被选中去淘汰的概率),即k/(k+1)*(1/k) = 1/(k+1)。
那么被留下的概率就是 1 - (被淘汰的概率) = 1- 1/(k+1) = k/(k+1)。
但是这个概率还要乘上原来被保留的概率k/k,也就是1啦,所以不用乘了,就是k/(k+1)
这样一来,所有的数据被留下的概率都是k/(k+1),满足等概率抽取,得证。
(3)推广到k+2,k+3到n和步骤二中是一样的道理。
假设现在增加了一个数,第m个数,用k/m的概率选择是否保留,此时第m个数被保留的概率是k/m。
对于蓄水池中的k个数,它们原来的被保留概率是k/(m-1),但是现在新来了一个数据,每个数据都面临被淘汰的风险。
淘汰的概率为(第m个数被选中的概率)乘以(每个数据被选中去淘汰的概率),即k/m * (1/k) = 1/m。
那么被留下来的概率就是1 - (被淘汰的概率)= 1 - 1/m = (m-1)/m。
但是这个概率还要乘上原来的被保留概率k/(m-1),即 (m-1)/m * ( k/(m-1)) = k/m。
这样一来,所有的数据被留下的概率都是k/m,满足等概率抽取,得证。
时间复杂度O(n),空间复杂度O(k),k为要取的数的个数,本题中k等于1。
3. 代码
每次random索引,然后遍历链表的做法
# 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):
self.head = head
p = head
i = 0
while p != None:
i += 1
p = p.next
self.n = i
def getRandom(self):
i = random.randint(0,self.n-1)
p = self.head
while i:
p = p.next
i -= 1
return p.val
先用数组存起来,每次random索引后直接取得
class Solution(object):
def __init__(self, head):
self.nums = []
while head != None:
self.nums.append(head.val)
head = head.next
def getRandom(self):
i = random.randint(0,len(self.nums)-1)
return self.nums[i]
蓄水池采样
import random
class ListNode(object):
def __init__(self, x):
self.val = x
self.next = None
class Solution(object):
def __init__(self, head):
self.head = head
def getRandom(self):
p = self.head
num = p.val
count = 2
while p.next:
p = p.next
if(random.random() < 1.0/count):
num = p.val
count += 1
return num