この記事は、面接後の会話記録をストリームとして扱う技術メモを、Reservoir Sampling の実装に絞って書き直したものです。背景の説明は面接トランスクリプトレビューの技術メモにあります。

TypeScriptでReservoir Samplingを実装する：長さ不明のストリームから公平にサンプリングする

ログ、イベント、文字起こしのように「最後に何件来るか分からない」データから、メモリを増やさずにランダムな k 件を取りたいことがあります。

このとき使えるのが Reservoir Sampling です。1 パスで処理でき、保持する要素は常に k 件だけです。しかも、入力の各要素が最後に選ばれる確率は同じになります。

この記事では、次の実装を TypeScript で作ります。

入力全体を配列にせず、O(k) メモリで処理する
AsyncIterable にそのまま接続する
乱数生成器を注入して、境界と公平性を検証する
「なぜ公平なのか」を確率で説明する

先に結論

入力を 1 件ずつ i = 1, 2, ... と数えます。

i <= k なら、そのまま reservoir に入れる
i > k なら、0 以上 i - 1 以下の整数 j を一様に選ぶ
j < k のときだけ、reservoir[j] を新しい要素で置き換える

j >= k なら、その要素は捨てます。

Math.random() は [0, 1) を返すので、Math.floor(Math.random() * i) で 0..i-1 を作れます。

なぜ配列に全件入れてから選んではいけないのか

全件を配列に入れてから k 件を選ぶ方法は簡単です。しかし、入力が大きいとメモリ使用量が O(n) になります。ストリームの終端が見えない場合は、そもそも「最後まで待つ」ことができません。

先頭から常に k 件だけを保存する方法も公平ではありません。後から来た要素が一度も選ばれないからです。

Reservoir Sampling は、後から来た要素にも置換の機会を与えます。保存領域を増やさずに、全要素を同じ条件で扱える点が重要です。

TypeScript実装

乱数を Rng として切り出します。テストでは乱数を固定でき、実運用ではデフォルトの Math.random を使えます。

type Rng = () => number;

export class Reservoir<T> {
  private readonly items: T[] = [];
  private seen = 0;

  constructor(
    private readonly capacity: number,
    private readonly rng: Rng = Math.random,
  ) {
    if (!Number.isInteger(capacity) || capacity <= 0) {
      throw new RangeError("capacity must be a positive integer");
    }
  }

  push(value: T): void {
    this.seen += 1;

    if (this.items.length < this.capacity) {
      this.items.push(value);
      return;
    }

    const u = this.rng();
    if (!(u >= 0 && u < 1)) {
      throw new RangeError("rng must return a value in [0, 1)");
    }

    // seen 件目なので、候補の添字は 0..seen-1。
    const index = Math.floor(u * this.seen);

    if (index < this.capacity) {
      this.items[index] = value;
    }
  }

  values(): readonly T[] {
    return this.items;
  }
}

seen は「いま何件目を読んだか」です。items.length ではありません。最初の k 件を入れた後は、保持数がずっと k のままだからです。

乱数の範囲も検証しています。乱数源を外から注入できる設計にしておくと、再現性のあるテストを書けます。

AsyncIterableに接続する

HTTP のページング、ログ購読、ストリーミング文字起こしなどは AsyncIterable にすると、サンプリング処理から入力の詳細を分離できます。

export async function sampleAsync<T>(
  source: AsyncIterable<T>,
  capacity: number,
  rng: Rng = Math.random,
): Promise<readonly T[]> {
  const reservoir = new Reservoir<T>(capacity, rng);

  for await (const value of source) {
    reservoir.push(value);
  }

  return reservoir.values();
}

使用例です。

async function* transcriptLines(): AsyncGenerator<string> {
  yield "質問: 最近担当した機能は？";
  yield "回答: ...";
  yield "追質問: p99 レイテンシは？";
  yield "回答: ...";
  yield "追質問: 失敗時の再試行は？";
}

const sample = await sampleAsync(transcriptLines(), 2);
console.log(sample);

入力は一度しか走査されません。transcriptLines が数百万行になっても、保持する行数は 2 件です。

公平性を確かめる

k = 1 のとき、i 件目が最後まで残る確率は次のようになります。

1 件目が残る確率：最初に保持され、2 件目で置換されない確率 1/2、3 件目でも置換されない確率 2/3、…
i 件目まで残る確率：1/i
したがって、最後の n 件を見終わったとき、どの要素も 1/n

k 件の場合も同じ考え方です。i 件目が reservoir に入る確率は k/i、その後の各要素で置換されない確率を掛け合わせると、最後まで残る確率はすべて k/n になります。

コードでも確認できます。5 件から 1 件を選ぶ試行を 20,000 回行うと、各要素はおよそ 4,000 回ずつ選ばれます。

const counts = [0, 0, 0, 0, 0];

for (let trial = 0; trial < 20_000; trial += 1) {
  let state = (trial + 1) | 0;

  const rng = () => {
    // テスト用の決定的な xorshift32
    state ^= state << 13;
    state ^= state >>> 17;
    state ^= state << 5;
    return ((state >>> 0) % 1_000_000) / 1_000_000;
  };

  const reservoir = new Reservoir<number>(1, rng);
  for (let value = 0; value < 5; value += 1) {
    reservoir.push(value);
  }
  counts[reservoir.values()[0]] += 1;
}

console.log(counts);
// 例: [4046, 3941, 3971, 4012, 4030]

この出力は「完全に同じ」を意味しません。乱数なので揺れます。ここで見るべきなのは、極端な偏りがないことです。乱数源を固定しておけば、CI でも同じ入力を再現できます。

実装で迷いやすい境界

capacity = 0 を許すか

この実装では拒否しています。0 件を返したい場合は呼び出し側で処理したほうが、j < capacity の意味が明確です。0 を許す API にすると、最初の入力を捨てる分岐と乱数を呼ぶ分岐が増えます。

Math.random() の品質

公平性の検証には十分ですが、暗号用途の乱数ではありません。抽選の監査やセキュリティ上の要件がある場合は、crypto.getRandomValues など、要件に合った乱数源を注入します。

重み付きサンプルとの違い

Reservoir Sampling が保証するのは「各要素を同じ確率で選ぶ」ことです。重要度やスコアに応じて選びたいなら、重み付き Reservoir Sampling という別のアルゴリズムが必要です。Math.random() の値を加工するだけでは、重み付きにはなりません。

並列処理

複数ワーカーが別々にサンプリングした結果を単純に連結すると、全体の一様性は壊れます。分割されたストリームを統合する場合は、各チャンクの件数を使って統合するアルゴリズムを選びます。まずは単一ストリームの 1 パス実装として使うのが安全です。

計算量

入力を n 件、保持数を k とすると、計算量は次の通りです。

項目	計算量
走査	O(n)
追加・置換	1 件あたり O(1)
追加メモリ	O(k)
入力の再走査	なし

「最後に何件あるか分からない」「全件を保存できない」という制約があるときに、候補を公平に絞り込めます。