大量のテキストからランダムに少数の行を抽出したい - Reservoir Sampling

前に以下のような記事を書きましたが、大量のテキストではうまくいかなかったので新たに書きました

ファイルからランダムにN行取り出す(shufコマンド) - 唯物是真 @Scaled_Wurm

上の記事ではテキストをランダムに\(k\)行取り出したい時"shuf -n k"コマンドでランダムにシャッフルした\(k\)行を取り出していました

ところが非常に大きなテキストファイルに対して上のコマンドを実行すると、一度にデータを全部メモリに読み込み始めているのか、すごい勢いでメモリを消費していきました(sort -Rでも)

そこでメモリをあまり使わずにランダムに\(k\)行取り出す方法について調べました

まず基本的な非復元抽出のアルゴリズムは以下の記事の発展手法とか追記のあたりの話がわかりやすいと思います

非復元抽出の高速かつ実装が簡単な方法を考える - 睡眠不足？！

この記事の話も一度全部の要素を読み込んでいるので、今回の問題では使えません
(テキストを一度全部読み込んで要素数を調べて、インデックスを抽出してからまたテキストを読み込んで取り出せばできると思いますが……)

そこで一度ファイルを読み込むだけで(メモリに全部保存しておかないでも)ランダムに抽出できる方法を調べました
(実用的には、各行について乱数が一定確率以下になったら抽出とかでいいような気がしますが……)

Reservoir Sampling

Reservoir Samplingという方法があるみたいです
ちょっと元論文までは読めていないのですが、以下の複数の記事を参考にしてまとめました

Reservoir samplingは各要素をそれぞれ\(1\)回ずつみるだけでよいアルゴリズムです(全体の要素数がわからなくてよい)
記憶容量は抽出する要素数、今回の場合は\(k\)個しか必要ありません
出力される順番は完全なランダムではないので注意

アルゴリズムとコード

アルゴリズムを簡単に説明します

まず要素が\(k\)個になるまで要素を配列に読み込みます
\(n=k+1\)個目以降の要素の場合、\(0\)から\(n-1\)までの整数の乱数を生成してそれが\(k\)より小さければそのインデックスの要素を今の要素(\(n\)番目の要素)で置き換えます

以下のPythonのコードを見たほうがわかりやすいかもしれません
iterableがkより短い場合StopIterationの例外が出ます

import random
def reservoir_sampling(iterable, k):
    it = iter(iterable)
    reservoir = [next(it) for i in xrange(k)]
    n = k
    for item in it:
        n += 1
        r = random.randint(0, n - 1)
        if r < k:
            reservoir[r] = item
    return reservoir

if __name__ == '__main__':
    print reservoir_sampling(xrange(1000000), 10)

何故これでうまくいくのか。

\(n=k+1\)以降の要素について考える
\(n\)番目の要素が配列の要素を置き換える確率が\(\frac{k}{n}\)
配列中のある特定の要素が\(n+1\)番目に置き換えられる確率は、置き換えが起こる確率\(\frac{k}{n+1}\)と\(k\)個の内その要素が選ばれる確率\(\frac{1}{k}\)の積なので\(\frac{k}{n+1}\times\frac{1}{k}=\frac{1}{n+1}\)
よって、ある特定の要素が\(n+1\)番目に置き換えられない確率は\(1-\frac{1}{n+1}=\frac{n}{n+1}\)
なので\(n\)番目に配列にあった要素が\(N\)番目まで残っている確率は\(\frac{n}{n+1} \times \frac{n+1}{n+2}\times\dots\times\frac{N-2}{N-1}\times\frac{N-1}{N}=\frac{n}{N}\)となる
以上より\(n\)番目の要素が配列の要素を置き換えた後\(N\)番目まで残っている確率は、配列の要素を置き換える確率と\(N\)番目まで残る確率の積なので\(\frac{k}{n}\times\frac{n}{N}=\frac{k}{N}\)
最初に配列に含まれている要素については、\(k\)番目に配列にあった要素が\(N\)番目まで残っている確率と考えればよいので、上で計算したように\(\frac{k}{N}\)
この結果からすべての要素は等しい確率で配列に含まれることがわかる