什么是水塘抽样?
水塘抽样是从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,这一概率为
。
因此每一个元素被选中的概率均为1/n。
当k
为任意值时,
对于第
个元素而言,设其被选中的概率为
。在当前共有
个元素被收到的情况下,该元素当且仅当其后的每一个元素都未被选中到当前位置(设元素j被选中的概率为
,则元素j被选中到当前位置的概率为
)的情况下会被选中,这一概率应为
。证明如下: