水塘抽样(Reservoir Sampling): 力扣398题目而来
参考链接: https://zhuanlan.zhihu.com/p/29178293
问题描述:当内存无法加载全部数据时,如何从包含未知大小的数据流中随机选取k个数据,并且要保证每个数据被抽取到的概率相等。
特殊且常用情况:k=1,假设数据量为N
第一个数
n
1
n_1
n1,我们选择保留,
p
(
n
1
)
=
1
p(n_1) = 1
p(n1)=1
第二个数
n
2
n_2
n2,我们以
1
2
\frac{1}{2}
21的概率保留,那么
n
1
n_1
n1被保留(不被覆盖)的概率
p
(
n
1
)
=
1
∗
(
1
−
1
2
)
=
1
2
p(n_1) = 1*(1-\frac{1}{2}) = \frac{1}{2}
p(n1)=1∗(1−21)=21,
n
2
n_2
n2被选中(覆盖前面的数)的概率
p
(
n
2
)
=
1
2
p(n_2) = \frac{1}{2}
p(n2)=21
第三个数
n
3
n_3
n3,我们以
1
3
\frac{1}{3}
31的概率保留,那么
n
1
n_1
n1被保留的概率为
p
(
n
1
)
=
1
∗
(
1
−
1
2
)
∗
(
1
−
1
3
)
=
1
3
p(n_1)=1*(1-\frac{1}{2})*(1-\frac{1}{3}) = \frac{1}{3}
p(n1)=1∗(1−21)∗(1−31)=31,
n
2
n_2
n2被保留概率为
p
(
n
2
)
=
1
∗
1
2
∗
(
1
−
1
3
)
=
1
3
p(n_2) = 1*\frac{1}{2}*(1-\frac{1}{3}) = \frac{1}{3}
p(n2)=1∗21∗(1−31)=31,
n
3
n_3
n3被保留概率为
p
(
n
3
)
=
1
3
p(n_3) = \frac{1}{3}
p(n3)=31
…
第i个数
n
i
n_i
ni,我们以
1
i
\frac{1}{i}
i1的概率保留,那么
n
1
,
n
2
,
.
.
.
,
n
(
i
−
1
)
n_1, n_2, ..., n_(i-1)
n1,n2,...,n(i−1)被保留的概率也为
1
i
\frac{1}{i}
i1
一般情况,随机抽取k个样本,k>1
按照上述逻辑
对于前k个数字,我们选择保留
p
(
n
1
)
=
p
(
n
2
)
=
.
.
.
=
p
(
n
k
)
=
1
p(n_1)=p(n_2)=...=p(n_k) = 1
p(n1)=p(n2)=...=p(nk)=1
对于第k+1个数字,我们以 k k + 1 \frac{k}{k+1} k+1k的概率保留,那么前k个数字中 n r n_r nr被保留的概率为 上 一 轮 p ( n r ) 被 保 留 的 概 率 ∗ ( n k + 1 被 丢 弃 的 概 率 + n k + 1 被 保 留 ∗ n r 此 轮 不 被 丢 弃 ) 上一轮p(n_r)被保留的概率*(n_{k+1}被丢弃的概率+n_{k+1}被保留*n_r此轮不被丢弃) 上一轮p(nr)被保留的概率∗(nk+1被丢弃的概率+nk+1被保留∗nr此轮不被丢弃) = 1 ∗ ( 1 k + 1 + k k + 1 ∗ k − 1 k ) = k k + 1 1 * (\frac{1}{k+1} + \frac{k}{k+1}*\frac{k-1}{k}) = \frac{k}{k+1} 1∗(k+11+k+1k∗kk−1)=k+1k。
对于第k+2个数字,我们以
k
k
+
2
的
概
率
保
留
\frac{k}{k+2}的概率保留
k+2k的概率保留,前k个数字
n
r
n_r
nr被保留的概率为
p
(
n
r
)
=
k
k
+
1
∗
(
2
k
+
2
+
(
k
k
+
2
∗
k
−
1
k
)
)
=
k
k
+
2
p(n_r) = \frac{k}{k+1} * (\frac{2}{k+2} + (\frac{k}{k+2} * \frac{k-1}{k})) = \frac{k}{k+2}
p(nr)=k+1k∗(k+22+(k+2k∗kk−1))=k+2k
…
对于第i个数字
n
i
n_i
ni,i>k,我们以
k
i
\frac{k}{i}
ik的概率保留,那么前k个数字中
n
r
n_r
nr被保留的概率
p
(
n
r
)
=
k
i
−
1
∗
(
i
−
k
i
+
(
k
i
∗
k
−
1
k
)
)
=
k
i
p(n_r) = \frac{k}{i-1} * (\frac{i-k}{i} + (\frac{k}{i}*\frac{k-1}{k})) = \frac{k}{i}
p(nr)=i−1k∗(ii−k+(ik∗kk−1))=ik
python的简单实现
import random
def random_sample_k(nums, k):
res = [0]*k
for i in range(k): # 初始化前k个数字
res[i] = nums[i]
for i in range(k, len(nums)):
random_k = random.randrange(i) # 在i中随机选择一个数字
if random_k < k: # 如果数字小于k,即实现 k/i 的概率来决定是否覆盖
res[random_k] = nums[i]