水塘抽样算法

水塘抽样算法

给你一个未知长度的链表，请你设计一个算法，只能遍历一次，随机地返回链表中的一个节点。

这里说的随机是均匀随机（uniform random），也就是说，如果有n个元素，每个元素被选中的概率都是1/n，不可以有统计意义上的偏差。

分析：这个问题的难点在于，随机选择是「动态」的，比如说你现在你有 5 个元素，你已经随机选取了其中的某个元素a作为结果，但是现在再给你一个新元素b，你应该留着a还是将b作为结果呢，以什么逻辑选择a和b呢，怎么证明你的选择方法在概率上是公平的呢？

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

看代码容易理解这个思路：

/* 返回链表中一个随机节点的值 */
int getRandom(ListNode head) {
    Random r = new Random();
    int i = 0, res = 0;
    ListNode p = head;
    // while 循环遍历链表
    while (p != null) {
        // 生成一个 [0, i) 之间的整数
        // 这个整数等于 0 的概率就是 1/i
        if (r.nextInt(++i) == 0) {
            res = p.val;
        }
        p = p.next;
    }
    return res;
}

对于概率算法，代码往往都是很浅显的，这种问题的关键在于证明，你的算法为什么是对的？为什么每次以1/i的概率更新结果就可以保证结果是平均随机（uniform random）？

证明：假设总共有n个元素，我们要的随机性无非就是每个元素被选择的概率都是1/n对吧，那么对于第i个元素，它被选择的概率就是：

第i个元素被选择的概率是1/i，第i+1次不被替换的概率是1 - 1/(i+1)，以此类推，相乘就是第i个元素最终被选中的概率，就是1/n。

因此，该算法的逻辑是正确的。

同理，如果要随机选择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);
        // 这个整数小于 k 的概率就是 k/i
        if (j < k) {
            res[j] = p.val;
        }
        p = p.next;
    }
    return res;
}

对于数学证明，和上面区别不大：

因为虽然每次更新选择的概率增大了k倍，但是选到具体第i个元素的概率还是要乘1/k，也就回到了上一个推导