← 中文资源

简体中文版本

蓄水池抽样:数组与流式输入

用蓄水池算法从数组或单遍可迭代对象中等概率抽取固定数量样本,随机源可注入以便测试。

阅读英文原文

功能说明

reservoirSample 从数组中等概率抽取固定数量的样本,reservoirStream 则遍历任意可迭代对象,在只扫描一次且不保存全部数据的前提下完成抽样。两者都接受可注入的随机函数,便于在单元测试中固定结果。适用于日志采样、审计抽检、大文件随机取行以及流式数据概览等场景,内存占用只与样本大小有关。

关键代码

蓄水池先放入前 k 个元素,之后对第 i 个元素以 k/(i+1) 的概率替换池中随机一项,从而保证每个元素最终入选概率相同。流式版本用生成器实现,逐个消费输入并维护已见数量,最后一次性产出池内容;每个新元素只调用一次随机数,避免额外的数组复制。参数校验会拒绝负数或非整数容量。

使用注意

注入的随机函数必须返回 [0,1) 区间的值,否则索引可能越界,Math.random 满足该约定。样本顺序不保留原始顺序,若需要稳定展示应在抽样后自行排序。算法只保证均匀抽样,不保证不重复业务主键,去重需在输入侧完成。对流式生成器调用方应尽快消费,避免长时间持有迭代器导致上游资源无法释放。

验证方式

用固定随机序列测试:容量为 1 时断言返回首个元素或按预期替换;容量大于数组长度时返回全部元素;空数组返回空数组。统计十万次抽样的频次,确认各元素接近均匀。流式版本用生成器输入并断言只遍历一次,同时测试容量为 0、负数以及随机函数返回 0 的边界行为。

代码示例

export function reservoirSample(items, size, random = Math.random) {
  if (!Number.isInteger(size) || size < 0) {
    throw new RangeError("size must be a non-negative integer");
  }
  if (size === 0) return [];
  const reservoir = items.slice(0, size);
  for (let index = size; index < items.length; index += 1) {
    const candidate = Math.floor(random() * (index + 1));
    if (candidate < size) reservoir[candidate] = items[index];
  }
  return reservoir;
}

export function* reservoirStream(iterable, size, random = Math.random) {
  if (!Number.isInteger(size) || size < 0) {
    throw new RangeError("size must be a non-negative integer");
  }
  if (size === 0) return;
  const reservoir = [];
  let seen = 0;
  for (const item of iterable) {
    if (seen < size) {
      reservoir.push(item);
    } else {
      const candidate = Math.floor(random() * (seen + 1));
      if (candidate < size) reservoir[candidate] = item;
    }
    seen += 1;
  }
  yield* reservoir;
}