蓄水池抽样算法

一、预备知识

Java 随机数生成的方法：

java.util.Random

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21

/* 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() 就可以。

但如果是长度未知的海量数据流呢？该如何实现等概率采样？

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

三、实现原理（从未知长度的海量数据随机采样K个元素）

3.1 举例说明：从未知流中随机选择一个元素(K = 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个元素。

四、代码

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18

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 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51

/* 蓄水池算法： 只对重复值采用蓄水池算法 出现次数 概率 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. 链表中随机选择节点

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41

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; } }