蓄水池采样算法

转载自我的博客

一、预备知识

Java 随机数生成的方法：

java.util.Random

/*
  Main.java 
*/
public class Main {

    public static void main(String[] args) {

        Random random = new Random(2048);

        // 生成 0 ~ 99 之间的整数
        System.out.println( random.nextInt(100));

        // 生成 0 ~ 1.0 之间的小数
        System.out.println( random.nextDouble());

        // 生成 布尔
        System.out.println( random.nextBoolean());

    }
}

二、从未知流中随机等可能抽取K个数采样的需求

从固定区间内随机采样数据十分简单，直接调用 random.nextInt() 就可以。
但如果是长度未知的海量数据流呢？该如何实现等概率采样？

蓄水池采样算法就是一种解决方案。

三、实现原理

3.1 举例说明：从未知流中随机选择一个元素的实现方法

• 当数据流中只有一个数据：
  • 直接返回该数据
• 当数据流中有两个数据：
  • D0，D1 中随机选择一个。 概率均为 1/2
• 当数据流中有三个数据
  • Step1 : 处理 D0, D1 时 先保留一个，概率分别为 1/2
  • Step2 :处理 D3 时, 1/3 的概率保留D3, 1 - 1/3 的概率保留 Step1 中的结果
• 递推下去

3.2 解析：假设流的长度只有3

数据 D1 被采样概率：（1/2）* (2/3) = 1/3
数据 D2 被采样概率：（1/2）*(2/3) = 1/3
数据 D3 被采样概率： 1/3

3.3 算法描述

• 先选取数据流中的前k个元素，保存在池子pool中；
• 从第j（k + 1 <= j <= n）个元素开始
  • 每次先以概率 p = k/j选择是否让第j个元素留下；
    • 若j被选中，则从A中随机选择一个元素并用该元素j替换它
    • 否则直接淘汰该元素；
• 重复步骤2直到结束，最后集合A中剩下的就是保证随机抽取的k个元素。

四、代码实现

    private int[] sampling(int K) {
        int[] pool = new int[K];
        
        // 前 K 个元素直接放入数组中
        for (int i = 0; i < K; i++) { 
            pool[i] = stream[i];
        }

        for (int i = K; i < N; i++) { // K + 1 个元素开始进行概率采样
            int r = random.nextInt(i + 1);
            // 这里其实就是k/j的体现
            if (r < K) {
                pool[r] = stream[i];
            }
        }

        return pool;
    }

leetcode 题目：

T1. 随机数索引

/*

 蓄水池算法：
    只对重复值采用蓄水池算法
    出现次数     概率
    1            1
    2            1/2
    3            1/3
    4            1/4
*/
class Solution {

    private int[] nums;
    public Solution(int[] nums) {
        this.nums = nums;
    }
    
    public int pick(int target) {
        int index = getIndex(nums, target);
        return index;
    }

    private int getIndex(int[] nums, int target){
        // 统计出现的次数
        int count = 0;
        int index = -1;

        Random random = new Random();

        for (int i = 0; i < nums.length; i++){

            if (target == nums[i]){
                
                if (index == -1){
                    index = i;
                    continue;
                }
                count += 1;

                int r = random.nextInt(count + 1);

                if (r == 0){
                    index = i;
                }
            
            }
        }

        return index;
    }
}

T2. 链表中随机选择节点

class Solution {

    /** @param head The linked list's head.
        Note that the head is guaranteed to be not null, so it contains at least one node. */
    private ListNode head;    
    public Solution(ListNode head) {
        this.head = head;
    }
    
    /** Returns a random node's value. */
    public int getRandom() {
        int count = 0;
        ListNode now = head;
        int ret = -1;
        Random random = new Random();

        while (now != null){
            
            count += 1;

            if (ret == -1){
                
                ret = now.val;
                now = now.next;
                continue;
            }
            int rm = random.nextInt(count );


            if (rm == 0){
                ret =  now.val;
            }

            now = now.next;
            

        }

        return ret;
    }
}