水塘抽样 Reservoir sampling

太过爱你忘了你带给我的痛 2022-06-03 05:53 152阅读 0赞

# 什么是水塘抽样？ #

水塘抽样是从n个元素中随机选取k个元素的算法。其中n可以是一个非常大的或者未知的数字。通常来说，水塘抽样算法用于n超过内存的容量或n是一个非常大的输入流的抽样。

## 从已知数目的`n`个元素中抽样`k`个 ##

算法步骤：  
1. 首先将`n`个元素中的前`k`个放入存储返回值的数组（记为`ret`)中。  
2. 对于从剩余的`n - k`个元素，从第`k+1`个元素开始，记录当前已考虑过的元素的`count`，然后从`0`到`count - 1`之间取随机数记为`rand`，如果`rand`小于`k`，则将`ret[rand - 1]`赋值为当前元素。  
3. 直到将`n`个元素全部循环一遍。

Java实现：

public class ReservoirSample {
        public int[] reservoirSample(int[] input, int k) {
            int n = input.length;
            int[] ret = new int[k];
            for (int  i = 0; i < n; i++) {
                if (i < k) {
                    ret[i] = input[i];
                } else {
                    int rand = new Random().nextInt(i);
                    if (rand < k) {
                        ret[rand] = input[i];
                    }
                }
            }
    
            return ret;
        }
    }

## 从未知长度的输入流中抽样k个元素 ##

如果我们需要抽样的输入为一个流输入，那么我们就无法事先知道这个流的长度。但考虑我们上面的算法，就会发现其实其也适用于流输入。

相似地，我们任然用一个`count`来记录当前遇到的元素的总数。当`count`小于`k`的时候，我们就直接将当前元素放入`ret`，否则我们仍取`0`到`count-1`之间的随机数`rand`，当`rand`小于`k`的时候，我们将其放入`ret[rand]`。

## 证明 ##

对于`k=1`的情况，很容易通过数学归纳法证明。

> 基本情况:  
> 当仅有一个元素被接收到时，其被选择的概率为：1/1 = 1.  
> 递推:  
> 假设有i个元素就被收到，每一个元素被选中的概率为1/i.  
> 对于第(i+1)个元素而言, 当且仅当rand(i+1) == 0这个元素才会被选中，概率为1/(i + 1).  
> 那么当接收到i+1个元素时，第i个元素被选中的概率为 (1/i) \* i/(i + 1) = 1/(i + 1).。

因此可得当共有n个元素被收到时，每一个元素被选中的概率都为 1/n。  
  
我们也可以这样理解这个问题：  
当共有n个元素出现，如果我们选中了第i个元素，对于所有从第 (i+1)到第n个元素，rand(j)必须不为0，这一概率为P=1i×ii\+1×i\+1i\+2×⋯×n−1n=1n  P = 1 i × i i + 1 × i + 1 i + 2 × ⋯ × n − 1 n = 1 n 。  
因此每一个元素被选中的概率均为1/n。

当`k`为任意值时，  
对于第n  n 个元素而言，设其被选中的概率为Pn=kn Pn=k n。在当前共有N  N 个元素被收到的情况下，该元素当且仅当其后的每一个元素都未被选中到当前位置（设元素j被选中的概率为Pj Pj，则元素j被选中到当前位置的概率为Pjk  P j k ）的情况下会被选中，这一概率应为kN  k N 。证明如下：

P=Pn×∏Nj=n\+1(1−Pjk)  P = P n × ∏ j = n + 1 N ( 1 − P j k )   
=kn×∏Nj=n\+1(1−kk×j)  = k n × ∏ j = n + 1 N ( 1 − k k × j )   
=kn×∏Nj=n\+1(j−1j)  = k n × ∏ j = n + 1 N ( j − 1 j )   
=kn×nn\+1×⋯×N−1N  = k n × n n + 1 × ⋯ × N − 1 N   
=kN  = k N