python 蓄水池采样代码蓄水池采样法

转载

huatechinfo 2023-11-24 00:38:53

文章标签 python 蓄水池采样代码数组数据随机数 文章分类 Python 后端开发

蓄水池采样算法

问题描述分析

采样问题经常会被遇到，比如：

从 100000 份调查报告中抽取 1000 份进行统计。
从一本很厚的电话簿中抽取 1000 人进行姓氏统计。
从 Google 搜索 "Ken Thompson"，从中抽取 100 个结果查看哪些是今年的。

这些都是很基本的采用问题。

既然说到采样问题，最重要的就是做到公平，也就是保证每个元素被采样到的概率是相同的。所以可以想到要想实现这样的算法，就需要掷骰子，也就是随机数算法。（这里就不具体讨论随机数算法了，假定我们有了一套很成熟的随机数算法了）

对于第一个问题，还是比较简单，通过算法生成 \([0, 100000 - 1)\)

但是对于第二和第三个问题，就有些不同了，我们不知道数据的整体规模有多大。可能有人会想到，我可以先对数据进行一次遍历，计算出数据的数量 \(N\)，然后再按照上述的方法进行采样即可。这当然可以，但是并不好，毕竟这可能需要花上很多时间。也可以尝试估算数据的规模，但是这样得到的采样数据分布可能并不平均。

算法过程

终于要讲到蓄水池采样算法（Reservoir Sampling）了。先说一下算法的过程：

假设数据序列的规模为 \(n\)，需要采样的数量的为 \(k\)。

首先构建一个可容纳 \(k\) 个元素的数组，将序列的前 \(k\)

然后从第 \(k+1\) 个元素开始，以 \(\frac{k}{n}\)

证明过程

对于第 \(i\) 个数（\(i \le k\)）。在 \(k\) 步之前，被选中的概率为 \(1\)。当走到第 \(k + 1\) 步时，被 \(k + 1\) 个元素替换的概率 = \(k + 1\) 个元素被选中的概率 * \(i\) 被选中替换的概率，即为 \(\frac{k}{k + 1} \times \frac{1}{k} = \frac{1}{k + 1}\)。则被保留的概率为 \(1 - \frac{1}{k + 1} = \frac{k}{k + 1}\)。依次类推，不被 \(k + 2\) 个元素替换的概率为 \(1 - \frac{k}{k + 2} \times \frac{1}{k} = \frac{k + 1}{k + 2}\)。则运行到第 \(n\)

\[1 \times \frac{k}{k + 1} \times \frac{k + 1}{k + 2} \times \frac{k + 2}{k + 3} \times … \times \frac{n - 1}{n} = \frac{k}{n} \]

对于第 \(j\) 个数（\(j > k\)）。在第 \(j\) 步被选中的概率为 \(\frac{k}{j}\)。不被 \(j + 1\) 个元素替换的概率为 \(1 - \frac{k}{j + 1} \times \frac{1}{k} = \frac{j}{j + 1}\)。则运行到第 \(n\)

\[\frac{k}{j} \times \frac{j}{j + 1} \times \frac{j + 1}{j + 2} \times \frac{j + 2}{j + 3} \times ... \times \frac{n - 1}{n} = \frac{k}{n} \]

所以对于其中每个元素，被保留的概率都为 \(\frac{k}{n}\).

代码示例

贴出测试用的示例代码（Java 实现）：

public class ReservoirSamplingTest {

    private int[] pool; // 所有数据
    private final int N = 100000; // 数据规模
    private Random random = new Random();

    @Before
    public void setUp() throws Exception {
        // 初始化
        pool = new int[N];
        for (int i = 0; i < N; i++) {
            pool[i] = i;
        }
    }

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

        for (int i = K; i < N; i++) { // K + 1 个元素开始进行概率采样
            int r = random.nextInt(i + 1);
            if (r < K) {
                result[r] = pool[i];
            }
        }

        return result;
    }

    @Test
    public void test() throws Exception {
        for (int i : sampling(100)) {
            System.out.println(i);
        }
    }
}

结果就不贴出来了，毕竟每次运行结果都不一样。