如果让你设计一个抽奖算法,从N个参与者中选出M个中奖者,你会怎么做?
假设这样的一个场景:给定一个数据流,数据流的长度N非常大,如何能在只遍历一遍数据的情况下,随机选出M个数据,并且使得N个数据被选中的概率全都相等呢?
从这个场景中我们可以得出这样这样几个结论:
- 数据流的长度N非常大,意味着我们不可能遍历后本地存储,因为我们肯定要使用数据结构保存抽奖后的结果,因此空间复杂度要求是O(M),当M远小于N时可以近似认为空间复杂度是O(1);
- 题目要求只遍历一遍数据,即时间复杂度要求O(N);
- 要求N个数据被选中的概率相等, 即每个小球被选中的概率都是M/N
这便是著名的“蓄水池算法”。
“蓄水池算法”的流程如下:
- 使用一个池子来保存中奖的数据,初始状态奖池为空;
- 遍历给到的数据流,假设遍历到第i个数据,则
- 如果i<=M,那么就把这个数据放入奖池,此时奖池中所有数据的中奖概率都是100%
- 如果 i > M,则先算出这个数据是否要进入奖池,概率是M/i
- 如果数据要进入奖池,再算出要把奖池中的哪个元素扔出奖池,每个元素被扔出的概率都是1/M。
我们假设M = 10,当遍历前10个元素时,每个元素都100%进入奖池;那么遍历到第11个元素时,这个元素进入奖池的概率是多少呢?
这个时候奖池中的元素肯定是1~10号元素,第11个元素进入奖池的概率就是10/11,目光集中到3号元素,根据条件概率的计算公式,11号元素遍历完后,3号元素被扔出奖池的概率是 11号元素进入奖池的概率 * 3号元素被扔出的概率,也就是10/11 * 1/10 = 1/11,那么此时3号元素仍然在奖池中的概率也是10/11;
遍历到第12号元素时,3号元素仍然不被扔出去的概率也是1 - 10/12 * 1/10 = 11/12,因此12号元素遍历完成后,3号元素仍在奖池里的概率是10/11 * 11/12 = 10/12;
遍历到13号元素时,3号元素仍不被扔出去的概率是1-10/13 * 1/10 = 12/13,因此13号元素遍历完成后,3号元素仍在奖池里的概率是10/11 * 11/12 * 12/13 = 10/13
以此类推,假设此时遍历到2023号元素,2023号元素进入奖池的概率是10/2023,仍然把目光集中在3号元素,此时3号元素仍在奖池中的概率是 10/11 * 11/12 * 12/13 * ... * 2022/2023 = 10/2022 * 2022/2023 = 10/2023。
以上就是蓄水池算法的整体思想。蓄水池算法一般被用于解决不确定数量数据的随机抽样算法,比如leetcode382题 链表随机节点 ,通过蓄水池算法可以用O(N)的时间复杂度解决,以及最近字节跳动一道的面试题,也可以用减配版本的蓄水池算法解决,我把题目贴在下面,感兴趣的小伙伴可以试一试。
给定一个字符串abcdabgh,给个字符a,随机返回a下标,比如这个是0 4。要求返回的概率必须一样,空间复杂度要求O1即不能开任何空间存储下标,并且只能遍历一次。