根据您想要的运行时保证,有一种著名的 O(n) 算法可以一次从数字流中挑选 k 个随机元素。为了理解这个算法,让我们先看一下我们想从集合中挑选一个元素的情况,然后我们将把它推广到挑选 k 个元素的情况下。这种方法的优点是它不需要任何关于输入集大小的预先知识,并保证对元素进行可证明的均匀采样,这总是非常好的。
假设我们想从集合中挑选一个元素。为此,我们将遍历集合中的所有元素,并且在每一点都将维护一个我们计划返回的候选元素。当我们遍历元素列表时,我们会以一定的概率更新我们的猜测,直到最后我们选择了具有统一概率的单个元素。在每一点,我们都将保持以下不变量:
看到k个元素后,当前选择前k个元素中的任何一个作为候选元素的概率为1/k。
如果我们在整个数组中保持这个不变量,那么在看到所有 n 个元素之后,它们中的每一个都有 1 / n 的机会成为候选元素。因此,候选元素已以均匀随机概率进行采样。
要了解算法是如何工作的,让我们考虑一下它必须做些什么来保持不变量。假设我们刚刚看到了第一个元素。为了保持上述不变量,我们必须以概率 1 选择它,因此我们将初始猜测的候选元素设置为第一个元素。
现在,当我们来到第二个元素时,我们需要保持每个元素以 1/2 的概率被选择的不变量,因为我们已经看到了两个元素。所以让我们假设我们以 1/2 的概率选择第二个元素。然后我们知道以下内容:
- 我们选择第二个元素的概率是 1/2。
- 我们选择第一个元素的概率是我们第一次选择它的概率 (1) 乘以我们不只选择第二个元素的概率 (1/2)。这也是 1/2。
所以在这一点上仍然保持不变!让我们看看当我们来到第三个元素时会发生什么。此时,我们需要确保以 1/3 的概率选择每个元素。好吧,假设我们以 1/3 的概率选择最后一个元素。然后我们知道
- 我们选择第三个元素的概率是 1/3。
- 我们选择前两个元素中的任何一个的概率是在前两个步骤之后选择它的概率 (1/2) 乘以我们没有选择第三个元素的概率 (2/3)。这可以达到 1/3。
再说一次,不变量成立!
这里的一般模式是这样的:在我们看到 k 个元素之后,每个元素都有 1/k 的机会被选中。当我们看到第 (k + 1) 个元素时,我们以 1 / (k + 1) 的概率选择它。这意味着它以 1 / (k + 1) 的概率被选择,并且它之前的所有元素被选择的概率等于我们在 (1 / k) 之前选择它并且没有选择 (k + 1) 的几率)st 个元素这次是 (k / (k + 1)),这给了这些元素每个被选择的概率 1 / (k + 1)。由于这在每一步都保持不变,因此我们得到了一个很棒的算法:
- 当你看到它时,选择第一个元素作为候选。
- 对于每个连续元素,用概率为 1 / k 的元素替换候选元素,其中 k 是到目前为止看到的元素数。
这在 O(n) 时间内运行,需要 O(1) 空间,并从数据流中返回一个均匀随机的元素。
现在,如果我们想从集合中挑选出k个元素,而不仅仅是一个,让我们看看如何扩展它以使其工作。这个想法与之前的算法非常相似(实际上最终成为更通用算法的一个特例)。我们不是维护一个候选者,而是维护 k 个不同的候选者,存储在我们编号为 1、2、...、k 的数组中。在每一点,我们都保持这个不变量:
在看到 m > k 个元素后,选择前 m 个元素中的任何一个的概率为 k / m。
如果我们扫描整个数组,这意味着当我们完成时,每个元素都有被选中的概率 k / n。由于我们选择了 k 个不同的元素,这意味着我们从数组中随机均匀地采样元素。
该算法与之前类似。首先,以概率 1 从集合中选择前 k 个元素。这意味着当我们看到 k 个元素时,它们中的任何一个被选中的概率是 1 = k / k 并且不变量成立。现在,归纳假设在 m 次迭代后不变量成立,并考虑第 (m + 1) 次迭代。选择一个介于 1 和 (m + 1) 之间的随机数,包括 1 和 (m + 1)。如果我们选择一个介于 1 和 k(含)之间的数字,则将该候选元素替换为下一个元素。否则,不要选择下一个元素。这意味着我们根据需要选择概率为 k / (m + 1) 的下一个元素。选择前 m 个元素中的任何一个的概率是它们之前被选择的概率 (k / m) 乘以我们没有选择包含该元素的槽的概率 (m / (m + 1)),这给出了根据需要选择 k / (m + 1) 的总概率。通过归纳,这证明了该算法完全均匀地随机抽取集合中的 k 个元素!
而且,运行时间是 O(n),它与集合的大小成正比,完全与你要选择的元素数量无关。它也只使用 O(k) 内存,并且对所存储元素的类型不做任何假设。
由于您尝试为 C++ 执行此操作,作为一种无耻的自我推销,我在我的个人网站上提供了该算法的实现(以 STL 算法编写)。随意使用它!
希望这可以帮助!