水塘抽样,相等概率随机取数(不定长数据流)
问题描述:设 n u m s nums nums 中有 k k k 个值为 t a r g e t target target 的元素,保证这 k k k 个值成为返回值的概率均为 1 k \frac{1}{k} k1
操作:遍历 n u m s nums nums,当我们第 i i i 次遇到值为 t a r g e t target target 的元素时,随机选取 [ 0 , i ) [0,i) [0,i) 内的一个整数,如果其等于 0 0 0 (范围内任何值都可),那么就将返回值置为该元素下标,否则不变。
证明:
P
(
第
i
次
遇
到
值
为
t
a
r
g
e
t
的
元
素
的
下
标
成
为
最
终
返
回
值
)
=
P
(
第
i
次
随
机
=
0
)
×
P
(
第
i
+
1
次
随
机
!
=
0
)
×
…
×
P
(
第
k
次
随
机
!
=
0
)
=
1
i
×
(
1
−
1
i
+
1
)
×
…
(
1
−
1
k
)
=
1
i
×
i
i
+
1
×
…
×
k
−
1
k
=
1
k
\begin{aligned} &P(第 i 次遇到值为 target 的元素的下标成为最终返回值) \\ &= P(第 i 次随机 = 0) \times P(第 i+1 次随机 != 0) \times \ldots \times P(第 k 次随机 != 0) \\ & = \frac{1}{i} \times (1 - \frac{1}{i+1}) \times \ldots (1 - \frac{1}{k}) \\ & = \frac{1}{i} \times \frac{i}{i+1} \times \ldots \times \frac{k-1}{k} \\ & = \frac{1}{k} \end{aligned}
P(第i次遇到值为target的元素的下标成为最终返回值)=P(第i次随机=0)×P(第i+1次随机!=0)×…×P(第k次随机!=0)=i1×(1−i+11)×…(1−k1)=i1×i+1i×…×kk−1=k1
class Solution:
def __init__(self, head: Optional[ListNode]):
self.root = head
def getRandom(self) -> int:
node, ans, i = self.root, None, 0
while node:
if not randint(0, i):
ans = node.val
node, i = node.next, i + 1
return ans