当我谈蓄水池抽样算法,我谈些什么。

前些日子翻完了村上春树的《当我谈跑步时我谈些什么》,很喜欢,即使我不是一个跑者,即使我读完也并没有心生去做一个跑者的志向,也许只是因为村上的文字中特有的那种吸引力吧。本文标题化用了村上这本书的书名,当然了,村上的书名本身也是化用的别人的,这是题外话。

下面回到主题:蓄水池抽样算法。

我第一次知道这个算法其实是在我第一次去B厂面试的时候,面试官问了我这个问题,虽然当时他并没有告诉我这个叫蓄水池抽样。

问题是这样的,假设有一个数据流源源不断来到,而你想从中等概率抽取K个,如何抽取?

最大的难点在于,样本空间不是固定的,而是像一条水流一样不断流入,并且有可能在某一时刻停止。

这就是蓄水池抽样算法要解决的问题。而用来解决这个问题的蓄水池抽样算法其实也和这个问题一样,并不复杂或高深。

命运的公平与残酷

因为想要抽取K个,所以对前K个到来的样本,都把它们放入幸运儿水池。可以想象,如果水流很小,只来了K个或者甚至连K个样本都没有,则将它们悉数入选,这符合题目要求。

对于后面来的第K+i(i=1,…)个样本,以概率 $ \frac K {K+i} $ 选择留下它。这对该样本来说是公平的,毫无疑问。

池子就那么大,我们只要K个,新人到来,旧人只好被踢出局,命运就是这么残酷。

剔掉哪一个呢,自然是不偏不倚,随机选取一个淘汰出局咯,也就是说,以概率$\frac 1 K$,选取一个之前的幸运儿,把它踢掉,是的,无情地踢掉。

以上就是算法的全部了。要证明这样选取确实是等概率的,也很简单,用一下数学归纳法,或者直接按乘法公式写出概率式子(会发现分子分母消除规律),很容易不是吗。

路总是不止一条

其实还有一种更为简单的做法:为每一个到来的样本赋予一个0到1之间均匀分布的随机数,让这个随机数代表这个样本参与到生死竞争中,选取前K大(或者前K小)的随机数所代表的样本作为最后的K个幸运儿。

这个想法同样也比较直观,但证明略微复杂一点,要用到一点积分,以及我所喜欢的gamma函数的亲戚beta函数。

假设有N个样本,我们证明:从这N个样本中选取指定的K个(例如,$X_1$, $X_2$, …, $X_K$),其概率等于$\frac 1 {N \choose K}$.

注意到,必然有一个$x \in [0,1]$, x 将这K个幸运儿与其他 N-K 个时运不济的家伙分割开来,形成一道天然的屏障,左手是$Y$,右手是$X$。

K个幸运儿之间并无大小之分,它们大于x的概率为
$$
p(X>x) = \prod_{i=1}^K {p(X_i > x) = {(1-x)}^K}
$$

N-K 个卢瑟儿也无所谓顺序,只知道它们小于x, 其概率为
$$
F(x) = \prod_{i=1}^{N-K} {p(Y_i \le x) = {x^{N-K}}}
$$

则能够达成上述安排的概率为
$$
\int_0^1 p(Y=x)P(X>x)dx \\
= \int_0^1 {(1-x)}^K dF(x) \\
= (N-K) \int_0^1 {(1-x)^K \times {x}^{N-K-1}dx} \\
$$

再结合beta函数跟gamma函数的血亲关系,以及gamma函数对阶乘的扩展,易得最后之结论。

这种方法,因为要维护一个容量为K的堆,效率上要比第一种低一些,但有得必有失,它也有其优点,并且,这是一个‘致命’的优点:适合于分布式的场景。

还有一个问题要问

蓄水池抽样能够公平地选取一个集合,这是没有问题的。问题是,序列也是公平的吗?

答案是并非如此。按照两种抽取方法抽出有序的K个样本,每种样本序列出现的概率并不相等。

例如,从{1,2,3}中选2个样本,只能选出{1,3},{1,2},{3,2}, 选不出{3,1},{2,1},{2,3}

即使一开始知道N的大小,并且随机选取一个起始点,也是如此,公平的集合,不公平的序列。

要想序列公平,只能随机抽。如果抽到重复的,就丢了继续,直到抽满K个。