目录

蓄水池采样算法(Reservoir Sampling Algorithm)是一种随机抽样算法,用于从一个大的、未知大小的数据流中随机选择固定数量的样本。这种算法特别适用于数据流太大,无法一次性存储在内存中的情况。

原理

蓄水池采样算法的基本思想是将数据流中的元素视为一个“蓄水池”,我们的目标是从这个蓄水池中随机选取一定数量的样本。算法的步骤如下:

  1. 初始化:创建一个大小为 k 的数组(或列表)作为蓄水池,并将前 k 个数据元素直接放入蓄水池中。
  2. 遍历数据流:从第 k+1 个元素开始,对数据流中的每个元素进行遍历。
  3. 随机替换:对于每个新元素,生成一个从 1 到当前元素索引 n 的随机数 rand。如果 rand 小于 k,则用当前元素替换蓄水池中索引为 rand-1 的元素。
  4. 重复步骤3:直到遍历完所有数据元素。
  5. 输出结果:最终,蓄水池中的 k 个元素就是从整个数据流中随机选取的样本。

算法特点

  • 随机性:每个元素被选中的概率是相等的,都是 k/n,其中 n 是数据流中的元素总数。
  • 空间效率:只需要存储 k 个元素,无论数据流有多大。
  • 时间效率:对于每个新元素,算法的时间复杂度是常数时间,因此整个算法的时间复杂度是 O(n)。

应用场景

蓄水池采样算法在以下场景中非常有用:

  • 大数据流:当数据流太大,无法一次性加载到内存中时。
  • 实时分析:需要对数据流进行实时抽样分析时。
  • 随机抽样:需要从大量数据中随机抽取样本进行统计分析或机器学习训练时。

示例代码(Python)

import random

def reservoir_sampling(stream, k):
    reservoir = []
    for i, item in enumerate(stream):
        if i < k:
            reservoir.append(item)
        else:
            j = random.randint(0, i)
            if j < k:
                reservoir[j] = item
    return reservoir

# 示例使用
stream = range(1000000)  # 假设我们有一个非常大的数据流
k = 10  # 我们想要的样本数量
sample = reservoir_sampling(stream, k)
print(sample)

这个示例展示了如何使用蓄水池采样算法从一个大的数据流中随机抽取 k 个样本。

当然,以下是一个使用 C++ 实现蓄水池采样算法的示例。这个示例假设我们有一个整数类型的数据流,并希望从中随机抽取 k 个样本。

#include <iostream>
#include <vector>
#include <algorithm>
#include <random>

template<typename T>
class ReservoirSampling {
public:
    ReservoirSampling(size_t k) : k_(k), samples_(k) {}

    void AddSample(const T& sample) {
        if (count_ < k_) {
            samples_[count_++] = sample;
        } else {
            size_t index = std::uniform_int_distribution<size_t>(0, count_ - 1)(rng_);
            if (index < k_) {
                samples_[index] = sample;
            }
        }
    }

    const std::vector<T>& GetSamples() const {
        return samples_;
    }

private:
    size_t k_;
    size_t count_ = 0;
    std::vector<T> samples_;
    std::mt19937 rng_;
};

int main() {
    const size_t k = 5; // 我们想要的样本数量
    const size_t streamSize = 100; // 数据流的大小

    ReservoirSampling<int> sampler(k);

    std::random_device rd;
    std::mt19937 gen(rd());

    for (size_t i = 0; i < streamSize; ++i) {
        int sample = std::uniform_int_distribution<int>(0, 99)(gen);
        sampler.AddSample(sample);
    }

    std::vector<int> samples = sampler.GetSamples();
    std::cout << "Sampled elements: ";
    for (int sample : samples) {
        std::cout << sample << " ";
    }
    std::cout << std::endl;

    return 0;
}

具体练习

382. 链表随机节点

扩展

  • 线程安全性:使用互斥锁(std::mutex)保护共享数据。
  • 并发操作:通过提供线程安全的API来支持并发访问。
  • 增加删除数据的功能:添加从池中删除数据的功能。
  • 支持插入和删除元素的时间复杂度为 O(logn) 并且保持元素有序
graph TD A[main函数] --> B[创建ReservoirSampling对象] A --> C[生成随机样本并通过多线程添加到ReservoirSampling对象中] A --> D[获取并打印样本] B --> E[ReservoirSampling对象初始化] C --> F[多个线程并发调用AddSample] D --> G[调用GetSamples方法获取样本] E --> H[互斥锁初始化] F --> I[互斥锁保护AddSample方法并插入有序样本] G --> J[互斥锁保护GetSamples方法]
#include <iostream>
#include <set>
#include <random>
#include <mutex>
#include <thread>

template<typename T>
class ReservoirSampling {
public:
    ReservoirSampling(size_t k) : k_(k) {}

    void AddSample(const T& sample) {
        std::lock_guard<std::mutex> lock(mtx_);
        if (samples_.size() < k_) {
            samples_.insert(sample);
        } else {
            size_t index = std::uniform_int_distribution<size_t>(0, count_ - 1)(rng_);
            if (index < k_) {
                auto it = samples_.begin();
                std::advance(it, index);
                samples_.erase(it);
                samples_.insert(sample);
            }
        }
        ++count_;
    }

    void RemoveSample(const T& sample) {
        std::lock_guard<std::mutex> lock(mtx_);
        samples_.erase(sample);
    }

    std::vector<T> GetSamples() {
        std::lock_guard<std::mutex> lock(mtx_);
        return std::vector<T>(samples_.begin(), samples_.end());
    }

private:
    size_t k_;
    size_t count_ = 0;
    std::set<T> samples_;
    std::mt19937 rng_{std::random_device{}()};
    std::mutex mtx_;
};

int main() {
    const size_t k = 5; // 我们想要的样本数量
    const size_t streamSize = 100; // 数据流的大小

    ReservoirSampling<int> sampler(k);

    std::random_device rd;
    std::mt19937 gen(rd());

    std::vector<std::thread> threads;

    // 向采样器添加数据
    for (size_t i = 0; i < streamSize; ++i) {
        threads.emplace_back([&sampler, &gen]() {
            int sample = std::uniform_int_distribution<int>(0, 99)(gen);
            sampler.AddSample(sample);
        });
    }

    for (auto& thread : threads) {
        thread.join();
    }

    std::vector<int> samples = sampler.GetSamples();
    std::cout << "Sampled elements: ";
    for (int sample : samples) {
        std::cout << sample << " ";
    }
    std::cout << std::endl;

    return 0;
}

https://mp.weixin.qq.com/s/T6Q98ooerkM45866HSN5ZQ