水塘抽样(解决k个数据均匀抽样问题)

关注

原创最新推荐文章于 2026-03-06 06:48:15 发布 2k 阅读 AI 写同款· GEO 优化 ›

面试题— —从很长的数据流等概率随机采样的两种实现方法从很长的数据流等概率随机采样 “给出一个数据流，这个数据流的长度很大或者未知。并且对该数据流中数据只能访问一次。请写出一个随机选择算法，使得数据流中所有数据被选中的概率相等。” 在搜索引擎等相关领域也会被问到哦假设数据流长度为n，那么我们希望数据流中的每个数据都是以 1/n 的概率返回。相关题目： 382. 链表随机节点 398. 随机数索引不合适的但容易想到的两种方法：直观的做法是生成一个 1-n 的随机数 i，代表我们想取第 i 个数据。但是这个前提是我们知道了数据流的长度。通常来说，数据阅读详情

水塘抽样

文章目录

水塘抽样

背景

给定一个数据流，但是数据流的长度N很大，且N直到处理完所有数据之前都不可知，如何在只遍历一遍数据（O(N)）的情况下，能够随机选取出这组数据的k个概率相等的均匀抽样。(包含未知大小数据流中随机选取k个数据，使得每个数据抽取到的概率相等)

（1）仅扫描数据一次。

（2）空间复杂度为O(K)。空间复杂度与整个数据量无关，只与抽样大小有关。

（3）扫描到数据的前n 个数据时（n>k）,保存当前已扫描数据的k个均匀抽样。

思路

将取出数据数量k当做一个水塘

（1）如果接受的数据量小于k，则依次放入采样数组中

（2）当接收到第i个数据，i大于等于k时，在[0,i]的范围内取一个随机数d 如果d落在了[0,k-1]的范围内，则取接收到的第i个数据替换采样数组中下标等于d位置上的值

情况一：初始化的时候，也就是水塘未满之前，每一个数被选中的概率都一样都是1
情况二：水塘满了的时候

这时候

（1）当k+1元素未被选中的时候，数组中没有元素替换。所以k+1个元素被选中的概率为 k/k+1

（2）当k+1元素被选中的时候，数组中某个元素被k+1个元素替换掉了。那么第k+1个元素被选中的概率为 $k / k + 1$ ，这就是新元素出现在水塘中的概率。水库中任意一个元素被替换掉的概率是：(k/k+1)*(1/k)=1/(k+1)，意即首先要第k+1个元素被选中，然后该元素在k个元素中被选中。那它未被替换的概率就是1-1/(k+1)=k/(k+1)。可以看出来，旧元素和新元素出现的概率是相等的。

k=1

首先考虑简单的情况，当k=1时，如何制定策略：

假设数据流含有N个数，我们知道如果要保证所有的数被抽到的概率相等，那么每个数抽到的概率应该为 [外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-nq0gHLrN-1650884011739)(https://www.zhihu.com/equation?tex=%5Cfrac%7B1%7D%7BN%7D)] 。

那我们可以这样做：

遇到第1个数 [外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-V1PvVcTC-1650884011741)(https://www.zhihu.com/equation?tex=n_1)] 的时候，我们保留它， [外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-V3dzQNBl-1650884011741)(https://www.zhihu.com/equation?tex=p%28n_1%29%3D1)]
遇到第2个数 [外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-r1guQco1-1650884011742)(https://www.zhihu.com/equation?tex=n_2)] 的时候，我们以 [外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-JiofxRR9-1650884011742)(https://www.zhihu.com/equation?tex=%5Cfrac%7B1%7D%7B2%7D)] 的概率保留它，那么 [外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-CdxGm8f4-1650884011742)(https://www.zhihu.com/equation?tex=p%28n_1%29%3D1%5Ctimes+%5Cfrac%7B1%7D%7B2%7D%3D%5Cfrac%7B1%7D%7B2%7D)] ，[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-3FsovM7w-1650884011743)(https://www.zhihu.com/equation?tex=p%28n_2%29%3D%5Cfrac%7B1%7D%7B2%7D)]
遇到第3个数 [外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-sLoKUVDI-1650884011743)(https://www.zhihu.com/equation?tex=n_3)] 的时候，我们以 [外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-xztwuuHd-1650884011743)(https://www.zhihu.com/equation?tex=%5Cfrac%7B1%7D%7B3%7D)] 的概率保留它，那么 [外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-88vDNQMU-1650884011744)(https://www.zhihu.com/equation?tex=p%28n_1%29%3Dp%28n_2%29%3D%5Cfrac%7B1%7D%7B2%7D%5Ctimes%281-%5Cfrac%7B1%7D%7B3%7D%29%3D%5Cfrac%7B1%7D%7B3%7D)] ， [外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-di4Yq5pz-1650884011744)(https://www.zhihu.com/equation?tex=p%28n_3%29%3D%5Cfrac%7B1%7D%7B3%7D)]
……
遇到第i个数 [外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-7uXePSG0-1650884011745)(https://www.zhihu.com/equation?tex=n_i)] 的时候，我们以 [外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-5YTToCRK-1650884011745)(https://www.zhihu.com/equation?tex=%5Cfrac%7B1%7D%7Bi%7D)] 的概率保留它，那么 [外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-MbfOdDFV-1650884011745)(https://www.zhihu.com/equation?tex=p%28n_1%29%3Dp%28n_2%29%3Dp%28n_3%29%3D%5Cdots%3Dp%28n_%7Bi-1%7D%29%3D%5Cfrac%7B1%7D%7Bi-1%7D%5Ctimes%281-%5Cfrac%7B1%7D%7Bi%7D%29%3D%5Cfrac%7B1%7D%7Bi%7D)] ， [外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-zvBHmOX8-1650884011746)(https://www.zhihu.com/equation?tex=p%28n_i%29%3D%5Cfrac%7B1%7D%7Bi%7D)]

这样就可以看出，对于k=1的情况，我们可以制定这样简单的抽样策略：

数据流中第i个数被保留的概率为 [外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-n6XovGX7-1650884011746)(https://www.zhihu.com/equation?tex=%5Cfrac%7B1%7D%7Bi%7D)] 。只要采取这种策略，只需要遍历一遍数据流就可以得到采样值，并且保证所有数被选取的概率均为 [外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-Pl51IJdA-1650884011747)(https://www.zhihu.com/equation?tex=%5Cfrac%7B1%7D%7BN%7D)] 。

k>1

对于k>1的情况，我们可以采用类似的思考策略：

仍然假设数据流中含有N个数，那么要保证所有的数被抽到的概率相等，每个数被选取的概率必然为 [外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-vE2zuhEX-1650884011747)(https://www.zhihu.com/equation?tex=%5Cfrac%7Bk%7D%7BN%7D)] 。

对于前k个数 [外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-w6zaleOX-1650884011748)(https://www.zhihu.com/equation?tex=n_1%2Cn_2%2C%5Cdots%2Cn_k)] ，我们保留下来，则 [外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-SlUKpFTJ-1650884011748)(https://www.zhihu.com/equation?tex=p%28n_1%29%3Dp%28n_2%29%3D%5Cdots%3Dp%28n_k%29%3D1)] （下面连等采用 [外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-5SieuP3s-1650884011748)(https://www.zhihu.com/equation?tex=p%28n_%7B1%3Ak%7D%29)] 的形式）
对于第k+1个数 [外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-jhuHIfeC-1650884011749)(https://www.zhihu.com/equation?tex=n_%7Bk%2B1%7D)] ，我们以 [外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-xJUoEUXJ-1650884011749)(https://www.zhihu.com/equation?tex=%5Cfrac%7Bk%7D%7Bk%2B1%7D)] 的概率保留它（这里只是指本次被保留下来），那么前k个数中的 [外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-jp9ZixoD-1650884011749)(https://www.zhihu.com/equation?tex=n_r%28r%5Cin%7B1%3Ak%7D%29)] 被保留的概率可以这样表示： [外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-3GQclxXp-1650884011750)(https://www.zhihu.com/equation?tex=p%28n_r%E8%A2%AB%E4%BF%9D%E7%95%99%29%3Dp%28%E4%B8%8A%E4%B8%80%E8%BD%AEn_r%E8%A2%AB%E4%BF%9D%E7%95%99%29%5Ctimes%28p%28n_%7Bk%2B1%7D%E8%A2%AB%E4%B8%A2%E5%BC%83%29%2Bp%28n_%7Bk%2B1%7D%E8%A2%AB%E4%BF%9D%E7%95%99%29%5Ctimes+p%28n_r%E6%9C%AA%E8%A2%AB%E6%9B%BF%E6%8D%A2%29%29)] ，即 [外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-Y9ZB1pK9-1650884011750)(https://www.zhihu.com/equation?tex=p_%7B1%3Ak%7D%3D%5Cfrac%7B1%7D%7Bk%2B1%7D%2B%5Cfrac%7Bk%7D%7Bk%2B1%7D%5Ctimes+%5Cfrac%7Bk-1%7D%7Bk%7D%3D%5Cfrac%7Bk%7D%7Bk%2B1%7D)]
对于第k+2个数 [外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-ZqVI6VI0-1650884011750)(https://www.zhihu.com/equation?tex=n_%7Bk%2B2%7D)] ，我们以 [外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-QnsfJcbL-1650884011751)(https://www.zhihu.com/equation?tex=%5Cfrac%7Bk%7D%7Bk%2B2%7D)] 的概率保留它（这里只是指本次被保留下来），那么前k个被保留下来的数中的 [外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-T7UjKQzH-1650884011751)(https://www.zhihu.com/equation?tex=n_r%28r%5Cin%7B1%3Ak%7D%29)] 被保留的概率为 [外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-g14lvWGE-1650884011751)(https://www.zhihu.com/equation?tex=p_%7B1%3Ak%7D%3D%5Cfrac%7Bk%7D%7Bk%2B1%7D%5Ctimes%28%5Cfrac%7B2%7D%7Bk%2B2%7D%2B%5Cfrac%7Bk%7D%7Bk%2B2%7D%5Ctimes+%5Cfrac%7Bk-1%7D%7Bk%7D%29%3D%5Cfrac%7Bk%7D%7Bk%2B1%7D%5Ctimes%5Cfrac%7Bk%2B1%7D%7Bk%2B2%7D%3D%5Cfrac%7Bk%7D%7Bk%2B2%7D)]
……
对于第i（i>k）个数 [外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-RG2WvTua-1650884011752)(https://www.zhihu.com/equation?tex=n_i)] ，我们以 [外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-emJJyWpN-1650884011752)(https://www.zhihu.com/equation?tex=%5Cfrac%7Bk%7D%7Bi%7D)] 的概率保留它，前i-1个数中的 [外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-t3TqJYnU-1650884011752)(https://www.zhihu.com/equation?tex=n_r%28r%5Cin%7B1%3Ai-1%7D%29)] 被保留的概率为 [外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-haVXyWgY-1650884011752)(https://www.zhihu.com/equation?tex=p_%7B1%3Ak%7D%3D%5Cfrac%7Bk%7D%7Bi-1%7D%5Ctimes+%28%5Cfrac%7Bi-k%7D%7Bi%7D%2B%5Cfrac%7Bk%7D%7Bi%7D%5Ctimes+%5Cfrac%7Bk-1%7D%7Bk%7D%29%3D%5Cfrac%7Bk%7D%7Bi-1%7D%5Ctimes%5Cfrac%7Bi-1%7D%7Bi%7D%3D%5Cfrac%7Bk%7D%7Bi%7D)]

这样，我们可以制订策略：

对于前k个数，我们全部保留，对于第i（i>k）个数，我们以 [外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-7xqE1x0v-1650884011753)(https://www.zhihu.com/equation?tex=%5Cfrac%7Bk%7D%7Bi%7D)] 的概率保留第i个数，并以 [外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-fpAQfVOm-1650884011753)(https://www.zhihu.com/equation?tex=%5Cfrac%7B1%7D%7Bk%7D)] 的概率与前面已选择的k个数中的任意一个替换。

例题

398. 随机数索引

难度中等214收藏分享切换为英文接收动态反馈

给你一个可能含有 重复元素 的整数数组 nums ，请你随机输出给定的目标数字 target 的索引。你可以假设给定的数字一定存在于数组中。

实现 Solution 类：

Solution(int[] nums) 用数组 nums 初始化对象。
int pick(int target) 从 nums 中选出一个满足 nums[i] == target 的随机索引 i 。如果存在多个有效的索引，则每个索引的返回概率应当相等。

示例：

输入
["Solution", "pick", "pick", "pick"]
[[[1, 2, 3, 3, 3]], [3], [1], [3]]
输出
[null, 4, 0, 2]

解释
Solution solution = new Solution([1, 2, 3, 3, 3]);
solution.pick(3); // 随机返回索引 2, 3 或者 4 之一。每个索引的返回概率应该相等。
solution.pick(1); // 返回 0 。因为只有 nums[0] 等于 1 。
solution.pick(3); // 随机返回索引 2, 3 或者 4 之一。每个索引的返回概率应该相等。

提示：

1 <= nums.length <= 2 * 104
-231 <= nums[i] <= 231 - 1
target 是 nums 中的一个整数
最多调用 pick 函数 104 次

class Solution {
    vector<int> &nums;
public:
    Solution(vector<int> &nums) : nums(nums) {}

    int pick(int target) {
        int ans;
        for (int i = 0, cnt = 0; i < nums.size(); ++i) {
            if (nums[i] == target) {
                ++cnt; // 第 cnt 次遇到 target
                if (rand() % cnt == 0) {
                    ans = i;
                }
            }
        }
        return ans;
    }
};

作者：LeetCode-Solution
链接：https://leetcode-cn.com/problems/random-pick-index/solution/sui-ji-shu-suo-yin-by-leetcode-solution-ofsq/

这个解法，其实就是相当于在target的情况下，随机选取到1个数的概率是一致的。

所以