随机抽取样本问题&蓄水池算法&按权重抽取问题

2020-07-20 本文已影响0人快乐的二叉树

面试被问到的一个问题：从N个样本中随机抽取m个样本，要求每个样本被抽取的概率一致。升级1：要求精准抽到m个；升级2：对每个样本添加权重，要求抽取概率按照权重分配。

基础问题

问题描述：从N个样本中随机抽取m个样本，要求每个样本被抽取的概率一致，求怎么样抽取？数据量为百万级。
看到这个问题，最先想到的方法是，依次遍历每个样本，以 $\frac{m}{n}$ 的概率抽中当前样本作为最后 $m$ 中的一个，具体操作可以是：
1、每遍历一个样本，生成一个 $0\sim N$ 之间的随机数 $x$ ，对比 $x$ 和 $m$ 的大小；
2、若 $x$ 大于 $m$ ，说明属于 $\frac{1-m}{n}$ 概率内，不抽；若 $x$ 小于等于 $m$ ，说明属于 $\frac{m}{n}$ 概率内，抽它；
3、直到所有样本遍历结束。

还可以从另一个角度证明这个算法的公平性，对每个抽中的样本来说，它应该是被抽中的第 $i$ 个样本，那么它被抽中的概率是：第一次就被抽中的概率+第一次没抽中第二次被抽中的概率+...+前 $m-1$ 次都没抽中最后一次抽中的概率，用式子表示就是：

$\begin{align} P_{样本被抽中}&=\frac{1}{n}+\frac{n-1}{n}*\frac{1}{n-1}+...+\frac{n-1}{n}*\frac{n-2}{n-1}*...*\frac{1}{n-m+1} \\&=\frac{1}{n}+\frac{1}{n}+...+\frac{1}{n} \\&=\frac{m}{n} \end{align}$

不考虑调用随机数生成函数的耗时的话，这样做还有个问题，那就是最后抽中的数不一定正好是m个，因为一次遍历只保证了每个样本等概率被抽中，没法保证抽到的样本量。这时又想到，在遍历过程中要是抽满了m个，就退出循环停止遍历，可是当遍历完都没有抽满m个该怎么办呢？选择再遍历一次的话复杂度会很高，也可能出现遍历了很多次都没抽满的情况。

升级问题1

问题描述：从N个样本中随机抽取m个样本，要求每个样本被抽取的概率一致，而且保证最后正好抽到m个数。

其实不算是升级问题，因为在上个问题中其实已经规定了要抽取m个，只是因为优先想到的解法出现了bug，所以不得不再重新思考。

解法1

蓄水池算法可以很好地解决这个问题，但这里先不介绍它，先介绍另一种同样能实现的方式：第 $i$ 个样本，被抽中的概率是 $\frac {m-k}{n-i+1}$ ， $k$ 是已经抽中的样本个数。
1、第一个样本以概率 $\frac{m}{n}$ 抽取就好；
2、若第一个样本没抽中，则第二个样本抽中概率为 $(1-\frac{m}{n})*\frac{m}{n-1}$ ；若第一个样本被抽中了，那么第二个样本抽中的概率为 $\frac{m}{n}*\frac{m-1}{n-1}$ ，两种情况加起来，第二个样本被抽中的概率为 $(1-\frac{m}{n})*\frac{m}{n-1}+\frac{m}{n}*\frac{m-1}{n-1}=\frac{m}{n}$ ；
3、后面的样本依次类推，抽中概率和当前样本序号 $i$ 和已经抽中的样本数 $k$ 有关，最后可以得到每个样本被抽中的概率都是 $\frac{m}{n}$ 。

这个算法能够保证每个样本被抽到的概率都为 $\frac{m}{n}$ ，并且最后抽到的样本为m个。关键在于，每遍历或抽到一个样本之后，都要对接下来抽取的概率做调整，当抽取的很快时，概率的分子项会变小，后面样本越来越难被抽到；当抽取的比较慢，概率分子项会变大，后面样本被抽到的概率也会变大。而且当抽满m个之后，后面样本被抽到的概率就为0了；当前面的遍历一直没抽满值， $N$ 中只剩下 $m$ 个样本时，每个样本被抽中的概率变为1，所以怎么样都能满足条件。

解法2

接下来再看蓄水池算法，该算法是针对从一个长度为N的序列中随机抽取不重复的m个数，保证每个数被抽取到的概率为 $\frac{m}{n}$ 这个问题而构建的，算法步骤为：
1、构建一个可放m个元素的蓄水池，将序列的前m个元素放入蓄水池中；
2、从第m+1个元素开始，以 $\frac{m}{n}$ 的概率来决定该元素是否被替换到池子中；
3、当遍历完所有元素之后，蓄水池中的就是随机挑选出的m个元素。

算法伪代码为：

for i= m+1 to N
    k=random(1, i);
    if( k < m)
        SWAP the kth value and ith value
end for

上述算法的证明：

对于蓄水池中的前m个样本，最开始被选中的概率为1，然后每个样本留到最后的概率=（m+1到n的遍历中，每次替换都抽不到自己的概率），写成公式是 $P=1*\frac{m}{m+1}*\frac{m+1}{m+2}*...*\frac{n-1}{n}=\frac{m}{n}$ ；
对于蓄水池之外的样本，从第m+1个开始，设序号为j，它们最终能被换到蓄水池中的概率=遍历到自己的时候被换进去的概率*被换进去之后不再被换出来的概率，写成公式是 $P=\frac{m}{j}*\frac{j}{j+1}*\frac{j+1}{j+2}*...*\frac{n-1}{n}=\frac{m}{n}$

因此，不论刚开始是在蓄水池内还是在外，最后留在蓄水池内的概率都是一样的，而且这个算法一定保证了能选出m个样本来，因为一开始就是基于替换的思路。

升级问题2

问题描述：从N个样本中随机抽取m个样本，要求每个样本被抽取的概率一致。在此基础上，为每个样本分配一个权重值w，范围为[1,k]，表示权值为k的样本被抽中的概率是权值为1的样本概率的k倍。

解法很简单，在上面解法1的步骤中添加一个权重概率就好了：第 $i$ 个样本，被抽中的概率是 $\frac{w_i}{\sum_{i=1}^{i=N}w_i}*\frac {m-k}{n-i+1}$ ， $k$ 是已经抽中的样本个数, $w_i$ 表示第 $i$ 个样本的权重。
1、第一个样本以概率 $\frac{w_1}{\sum_{i=1}^{i=N}w_i}*\frac{m}{n}$ 抽取就好；
2、若第一个样本没抽中，则第二个样本抽中概率为 $(1-\frac{m}{n})*\frac{w_2}{\sum_{i=1}^{i=N}w_i}*\frac{m}{n-1}$ ；若第一个样本被抽中了，那么第二个样本抽中的概率为 $\frac{m}{n}*\frac{w_2}{\sum_{i=1}^{i=N}w_i}*\frac{m-1}{n-1}$ ，两种情况加起来，第二个样本被抽中的概率为 $(1-\frac{m}{n})*\frac{w_2}{\sum_{i=1}^{i=N}w_i}*\frac{m}{n-1}+\frac{m}{n}*\frac{w_2}{\sum_{i=1}^{i=N}w_i}*\frac{m-1}{n-1}=\frac{w_2}{\sum_{i=1}^{i=N}w_i}*\frac{m}{n}$ ；
3、后面的样本依次类推，抽中概率和当前样本序号 $i$ 和已经抽中的样本数 $k$ ，以及当前样本权重有关，最后可以得到每个样本被抽中的概率都是 $\frac{w_i}{\sum_{i=1}^{i=N}w_i}*\frac{m}{n}$ 。