如何等概率的返回已经遍历到的节点?

蓄水池算法

问题来源

1. 比如一个很大的文本,事先是不知道这个文本有多少行的,让我们随机返回其中一行,如何去做到随机的返回其中一行呢
2. 上述问题的隐藏含义:只能顺序读取每一行,且不能将整个文件加载到内存

例题1:

https://leetcode.cn/problems/linked-list-random-node/

给你一个单链表，随机选择链表的一个节点，并返回相应的节点值。每个节点 被选中的概率一样 。

实现 Solution 类：

Solution(ListNode head) 使用整数数组初始化对象。
int getRandom() 从链表中随机选择一个节点并返回该节点的值。链表中所有节点被选中的概率相等。

思路:

假设我们已经顺序遍历到了第i个节点,那应当保证这前i个节点的其中一个会被随机返回,也就是说这i个元素中的其中一个被选中的概率应当符合1/i

此外每次往后遍历时都应该保证这个概率.如果还想继续保证这个元素不被更新,所以有如下证明过程:

此处简单翻译一下这个公式:

也就是说读到第i个节点并想将其作为最终结果返回需满足上述公式推导,1/i就是将其选中,后续的因子都是在保护它不被新遍历到的元素更新.而0到i-1这么多个元素它爱咋选咋选,毕竟我第i个元素会把之前的选择的结果更新成当前的,所以也就无所谓了.1/n也表明有n个元素时,确实每个元素被返回的概率应当满足1/n的要求,这也就证明了蓄水池算法的正确性.

那如何将其翻译成程序呢?:

int i = 0;
int res = 0;
ListNode p = head;
while(p!=null){
    i++;
    if(0==random.nextInt(i)){
        res = p.val;
    }
    p=p.next;
}
return res;

解释:

[0,i)中共有i个数字可选,所以随机选中其中一个的概率就是1/i,那我们就用选0来达到目的吧(只要符合概率就行,又由于i==1的时候,只能选0,所以我们用选0来标定概率),简单推导,即可发现这样去写代码,是符合遍历一遍就能随机拿到其中一个节点值了.且最终(遍历最后一个节点的时候)每个节点的值被拿到的概率是1/n

完整代码:

class Solution {
    private ListNode head;
    public Solution(ListNode head) {
        //
        this.head = head;
    }

    public int getRandom() {
        Random r = new Random();
        int i=0;
        int res = 0;
        ListNode p = head;
        while(p!=null){
            i++;
            //生成[0,i)中的0的概率就是1/i
            if(0==r.nextInt(i)){
                res = p.val;
            }
            p=p.next;
        }
        return res;
    }
}

---

例题1的升级

如果要选择k个数,那只要在遍历到第i个元素时有k/i的概率选择该元素,以1-k/i的概率保持原有原则即可.那遍历到最后一个元素时,不就是每个节点都是k/n的概率被放在k个数中进行返回了吗.

翻译成代码:

int[] getRandom(ListNode head,int k){
    Random r = new Random();
    int[] res = new int[k];
    ListNode p = head;
    
    //将前K个元素先选上,因为前k个元素在遍历过程中100%会被选上呀
    for(int j = 0;j < k;j++){
        res[j]=p.val;
        p=p.next;
    }
    
    int i = k;
    while(p!=null){
        //生成一个[0,i)的随机整数
        //第一次就是[0,k]
        int j = r.nextInt(++i);
        //这个整数小于k的概率就是k/i,选中之后就将前k个元素的第j个元素替换成当前元素的值
        if(j<k){
            res[j]=p.val;
        }
        p=p.next;
    }
    return res;
}

证明:

当可选择的元素的个数本身就小于等于k个时,那这些个元素是100%会被选择的.这就不讨论了,所以当可选择的元素个数大于k个时:

例题2

https://leetcode.cn/problems/random-pick-index/

题干:

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

实现 Solution 类：

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

---

直接上代码吧,思路和上一题完全一样,只要能保证重复的元素(他们都等于target相等)的索引都能等概率被返回就可以解决问题

class Solution {
    private int[] nums;
    Random random = new Random();
    public Solution(int[] nums) {
        this.nums=nums;
    }

    public int pick(int target) {
        //题目的意思就是target必在数组当中,那target存在多份时,这些个元素的索引应当等概率的返回其中一个
        int res = 0;
        int i = 0;
        for(int k = 0;k<nums.length;k++){
            if(nums[k]==target){
                i++;
                if(0==random.nextInt(i)){
                    res = k;
                }
            }
        }
        return res;
    }
}

自己的理解,希望对读者有所帮助