LeetCode382之水塘抽样算法(相关话题：随机算法)

原创

数据与后端架构提升之路 2021-07-06 15:45:19 ©著作权

©著作权归作者所有：来自51CTO博客作者数据与后端架构提升之路的原创作品，请联系作者获取转载授权，否则将追究法律责任

问题描述

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

示例:

// 初始化一个单链表 [1,2,3].
ListNode head = new ListNode(1);
head.next = new ListNode(2);
head.next.next = new ListNode(3);
Solution solution = new Solution(head);

// getRandom()方法应随机返回1,2,3中的一个，保证每个元素被返回的概率相等。
solution.getRandom();

思路分析

先说结论，当你遇到第 i 个元素时，应该有 1/i 的概率选择该元素，1 - 1/i 的概率保持原有的选择。

举例说明： 1 - 10

遇到1，概率为1，保留第一个数。
遇到2，概率为1/2，这个时候，1和2各1/2的概率被保留
遇到3，3被保留的概率为1/3，(之前剩下的数假设1被保留)，2/3的概率 1 被保留，(此时1被保留的总概率为 2/3 * 1/2 = 1/3)
遇到4，4被保留的概率为1/4，(之前剩下的数假设1被保留)，3/4的概率 1 被保留，(此时1被保留的总概率为 3/4 * 2/3 * 1/2 = 1/4)
以此类推，每个数被保留的概率都是1/N。

下面是官话证明（个人觉得不好理解）

LeetCode382之水塘抽样算法(相关话题：随机算法)_算法

代码实现

用i来记录节点的个数，每次执行if(r.nextInt(i)==0)的判断就可以做到第i个节点被选中额概率为1/i

/**
 * Definition for singly-linked list.
 * public class ListNode {
 *     int val;
 *     ListNode next;
 *     ListNode() {}
 *     ListNode(int val) { this.val = val; }
 *     ListNode(int val, ListNode next) { this.val = val; this.next = next; }
 * }
 */
class Solution {

   private  ListNode head;


    public Solution(ListNode head) {
           this.head = head ;
    }
    
    /** Returns a random node's value. */
    public int getRandom() {

        int res=0,i=1;
        
        Random r = new Random();

        ListNode p = head;
        // while 循环遍历链表
        while (p != null) {
               // 生成一个 [0, i) 之间的整数
               // 这个整数等于 0 的概率就是 1/i
               if(r.nextInt(i)==0){
                 res =p.val;
               }
               i++;
               p=p.next;

        }


         return res;
    }
}

问题拓展

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

结论：如果要随机选择 k 个数，只要在第 i 个元素处以 k/i 的概率选择该元素，以 1 - k/i 的概率保持原有选择即可。

数学归纳法证明

代码：

/* 返回链表中 k 个随机节点的值 */
int[] getRandom(ListNode head, int k) {
    Random r = new Random();
    int[] res = new int[k];
    ListNode p = head;

    // 前 k 个元素先默认选上
    for (int j = 0; j < k && p != null; j++) {
        res[j] = p.val;
        p = p.next;
    }

    int i = k;
    // while 循环遍历链表
    while (p != null) {
        // 生成一个 [0, i) 之间的整数
        int j = r.nextInt(i);
        i++;
        // 这个整数小于 k 的概率就是 k/i
        if (j < k) {
            res[j] = p.val;
        }
        p = p.next;
    }
    return res;
}