蓄水池采样(Reservoir Sampling)算法
字数 1079 2025-11-12 16:41:47
蓄水池采样(Reservoir Sampling)算法
问题描述
蓄水池采样是一种在线随机采样算法,用于从包含未知大小或非常大的数据流中随机选取k个样本,使得每个样本被选中的概率相等。该算法只需要单次遍历数据,且内存使用量仅为O(k),与数据流大小无关。
核心思想
维护一个大小为k的"蓄水池"(数组),通过巧妙的概率替换策略,保证在处理完所有数据后,每个元素出现在蓄水池中的概率都是k/n(n为数据流总大小)。
算法步骤详解
步骤1:初始化蓄水池
- 创建一个大小为k的数组reservoir
- 读取数据流的前k个元素,直接放入蓄水池中
- 此时前k个元素被选中的概率为100%(因为蓄水池还未满)
步骤2:处理后续元素
- 从第k+1个元素开始(索引从0开始计,即第k个元素对应索引k)
- 对于第i个元素(i ≥ k):
a. 随机生成一个范围在[0, i)的整数j
b. 如果j < k,则用当前元素替换reservoir[j]位置的元素
c. 如果j ≥ k,则忽略当前元素
步骤3:概率分析
关键点在于理解为什么这样能保证等概率采样:
-
前k个元素的最终概率:
- 元素在初始时被选中概率 = 1
- 不被后面第i个元素替换的概率 = 1 - 1/(i+1) = i/(i+1)
- 因此最终概率 = 1 × (k/k+1) × (k+1/k+2) × ... × (n-1/n) = k/n
-
第i个元素(i > k)的选中概率:
- 被选入蓄水池的概率 = k/i
- 不被后面元素替换的概率 = (i/i+1) × (i+1/i+2) × ... × (n-1/n) = i/n
- 最终概率 = (k/i) × (i/n) = k/n
具体示例
假设k=3,数据流为[10, 20, 30, 40, 50],n=5
- 初始化:reservoir = [10, 20, 30]
- 处理第3个元素(索引3,值40):
- 生成随机数j ∈ [0,3)
- 如果j=0,1,2中任意一个,用40替换对应位置
- 处理第4个元素(索引4,值50):
- 生成随机数j ∈ [0,4)
- 如果j=0,1,2中任意一个,用50替换对应位置
最终每个元素在蓄水池中的概率都是3/5。
算法实现
import random
def reservoir_sampling(stream, k):
reservoir = []
# 步骤1:填充前k个元素
for i in range(k):
reservoir.append(stream[i])
# 步骤2:处理剩余元素
for i in range(k, len(stream)):
# 随机生成[0, i]的整数
j = random.randint(0, i)
# 如果落在前k范围内,替换对应元素
if j < k:
reservoir[j] = stream[i]
return reservoir
应用场景
- 大数据流中的随机采样
- 数据库查询结果的随机抽样
- 推荐系统的A/B测试分组
- 网络流量分析中的数据包采样
算法优势
- 单次遍历,时间复杂度O(n)
- 内存占用仅O(k),与数据流大小无关
- 在线算法,无需预先知道数据流大小
这种算法完美解决了大数据环境下无法全部加载到内存的随机采样问题。