Mersenne Twister でランダムに行をシャッフル

テキストファイルでランダムに行をシャッフルしたい場合,sort -R やその改良 (以下,sort -n) は標準入力を読めたり,メモリ周りを気にしなくて良かったりと使い勝手は非常に良いのだけど,Vapnik の原理が頭にあると sort が O(n log n) なのが少し気になるので,O(n) な Fisher-Yates shuffle (Durstenfeld's version) ベースの std::random_shuffle と std::tr1::mt19937 (または random ())を使って行をシャッフルしてみた.
Mersenne Twister でランダムに行をシャッフル (2) - ny23の日記 に sort -R|-n の代替になりうる実用版あり.
注意: MAX_MEM_SIZE=0 にすると,Random Write を行数回行うことになるので SSD上で実行しない方がいいかも

// mt-shuf.cc
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <vector>
#include <algorithm>

#ifdef USE_MT
#include <tr1/random>
#else
struct rand_ { long operator() (long max) const { return random () % max; } };
#endif

#ifndef MAX_MEM_SIZE // maximum data size to load on memory
#define MAX_MEM_SIZE (1UL << 34)
#endif

#ifndef MAX_LINE_LEN // maximum length of lines
#define MAX_LINE_LEN 65536
#endif

#ifndef TMP_FILE
#define TMP_FILE "/tmp/shuffle.tmp"
#endif

int main (int argc, char** argv) {
  if (argc < 2)
    { std::fprintf (stderr, "Usage: %s in [out]\n", argv[0]); std::exit (1); }
  // randomizer
#ifdef USE_MT
  std::tr1::mt19937 mt;
  std::tr1::uniform_int <> uni;
  std::tr1::variate_generator
    <std::tr1::mt19937, std::tr1::uniform_int <> > gen (mt, uni);
#else
  rand_ gen;
#endif
  // check size
  char * fn = argv[1];
  FILE * reader = std::fopen (fn, "rb");
  if (! reader)
    { std::fprintf (stderr, "no such file: %s\n", fn); std::exit (1);}
  if (std::fseek (reader, 0, SEEK_END) != 0) return -1;
  size_t size = std::ftell (reader);
  if (std::fseek (reader, 0, SEEK_SET) != 0) return -1;
  // std::fprintf (stderr, "data size: %ld bytes\n", size);
  
  // data on memory (if possible)
  char * buf = 0;
  if (size < MAX_MEM_SIZE) {
    // std::fprintf (stderr, "loading data onto memory..");
    buf = new char[size + 1];
    std::fread (buf, sizeof (char), size, reader);
    if (buf[size-1] != '\n') {
      std::fprintf (stderr, "WARNING: line feeder is appended\n");
      buf[size] = '\n'; ++size;
    }
  }
  std::fclose (reader);

  // shuffle
  FILE * lnfp = 0;
  char line[MAX_LINE_LEN];
  std::vector <size_t> lns;
  size_t ln = 0;
  if (buf) {
    for (size_t offset = 0; offset < size; ++offset) {
      lns.push_back (offset);
      ++ln;
      while (buf[offset] != '\n') ++offset;
      buf[offset] = '\0';
    }
    // std::fprintf (stderr, "done.\n");
  } else { // process on disk
    // std::fprintf (stderr, " Too large input; running as memory-saving mode.\n");
    // std::fprintf (stderr, "reading data..");
    reader = std::fopen (fn, "r");
    if (! MAX_MEM_SIZE) lnfp = std::fopen (TMP_FILE, "wb");
    size_t offset = 0;
    while ((std::fgets (&line[0], MAX_LINE_LEN, reader)) != NULL) {
      ++ln;
      if (MAX_MEM_SIZE)
        lns.push_back (offset);
      else
        std::fwrite (&offset, sizeof (size_t), 1, lnfp);
      size_t read = std::strlen (&line[0]);
      if (! std::feof (reader) && line[read-1] != '\n') {
        std::fprintf (stderr, "ERROR: Buffer Overflow; increase MAX_LINE_NUM\n");
        std::exit (1);
      }
      offset += read;
    }
    std::fclose (reader);
    if (! MAX_MEM_SIZE) std::fclose (lnfp);
    // std::fprintf (stderr, "done.\n");
  }
  // std::fprintf (stderr, "# lines: %ld\n", ln);
  // std::fprintf (stderr, "shuffling lines..");
  if (MAX_MEM_SIZE)
    std::random_shuffle (lns.begin (), lns.end (), gen);
  else { // too heavy random read/write
    lnfp = std::fopen (TMP_FILE, "rb+");
    for (size_t i = ln - 1; i > 0; --i) { //  Durstenfeld's algorithm on Disk
      size_t j = gen (i+1);
      size_t offset_i (0), offset_j (0);
      // read
      std::fseek (lnfp, i * sizeof (size_t), SEEK_SET);
      std::fread (&offset_i, sizeof (size_t), 1, lnfp);
      std::fseek (lnfp, j * sizeof (size_t), SEEK_SET);
      std::fread (&offset_j, sizeof (size_t), 1, lnfp);
      // swap & write
      std::fseek (lnfp, i * sizeof (size_t), SEEK_SET);
      std::fwrite (&offset_j, sizeof (size_t), 1, lnfp);
      std::fseek (lnfp, j * sizeof (size_t), SEEK_SET);
      std::fwrite (&offset_i, sizeof (size_t), 1, lnfp);
    }
    std::fclose (lnfp);
  }
  // std::fprintf (stderr, "done.\n");

  // save
  FILE * writer = argc < 3 ? stdout : std::fopen (argv[2], "w");
  // std::fprintf (stderr, "saving data..");
  if (buf) {
    for (std::vector <size_t>::iterator it = lns.begin ();
         it != lns.end (); ++it)
      std::fprintf (writer, "%s\n", &buf[*it]);
    delete [] buf;
  } else {
    reader = std::fopen (fn, "r");
    if (MAX_MEM_SIZE) {
      for (std::vector <size_t>::iterator it = lns.begin ();
           it != lns.end (); ++it) {
        std::fseek (reader, *it, SEEK_SET);
        if ((std::fgets (&line[0], MAX_LINE_LEN, reader)) != NULL) {
          size_t read = std::strlen (line);
          std::fwrite (line, sizeof (char), read, writer);
          if (line[read-1] != '\n') {
            std::fprintf (stderr, "WARNING: line feeder is appended\n");
            std::fprintf (writer, "\n");
          }
        }
      }
    } else {
      lnfp = std::fopen (TMP_FILE, "rb");
      size_t i = 0;
      while (std::fread (&i, sizeof (size_t), 1, lnfp)) {
        if (std::fseek (reader, i, SEEK_SET) != 0) return -1;
        if ((std::fgets (&line[0], MAX_LINE_LEN, reader)) != NULL) {
          size_t read = std::strlen (line);
          std::fwrite (line, sizeof (char), read, writer);
          if (line[read-1] != '\n') {
            std::fprintf (stderr, "WARNING: line feeder is appended\n");
            std::fprintf (writer, "\n");
          }
        }
      }
      std::fclose (lnfp);
    }
    std::fclose (reader);
  }
  std::fclose (writer);
  // std::fprintf (stderr, "done.\n");
};

というわけで実験してみる.MAX_MEM_SIZE=0 のときは in-place アルゴリズムとして動作するようにしてみた(やっつけで作ったので実装は非常に適当; あまり使わない方が良さそう; 08/11 バグ修正).

N=#lines (D=size) |  sort -R | sort -n | mt-shuf1| mt-shuf0 | mt-shuf | GNU shuf |
---------------------------------------------------------------------------------|
10^5 (    2.35MB) |    0.96s |   0.22s |   0.36s |    1.70s |   0.03s |    0.03s |
10^6 (   23.35MB) |   12.43s |   2.69s |   4.39s |   18.53s |   0.34s |    0.46s |
10^7 (  234.08MB) |  146.43s |  30.98s |  48.81s |  203.43s |   4.30s |    5.94s |
10^8 (2,341.42MB) | 1750.84s | 385.07s | 500.05s |      TBA |  51.02s |       NA |

sort -R: cat -n file | sort -R | cut -f 2- > /dev/ull
mt-shuf1 (64bit; MAX_MEM_SIZE=1; only line numbers on memory)
mt-shuf0 (64bit; MAX_MEM_SIZE=0; data on disk)
mt-shuf  (64bit; MAX_MEM_SIZE=1<<34; both data and line numbers on memory) 

mt-shuf, mt-shuf1, mt-shuf0 はそれぞれ,sizeof (size_t) * N + D bytes, sizeof (size_t) * N, 0 bytes の作業メモリを消費する.shuf の消費メモリは 32bit版 mt-shuffle より少し多い程度.mt-shuf* はディスクへのランダムアクセスが発生するため無残な結果に.一方,mt-shuf はメモリのみへのランダムアクセスになるので,サイズが大きくなるにつれキャッシュが効きにくくなって sort との計算量の違いがあまり出ていない感じ.数百M程度のデータなら使いではありそうだけど,実用から言えば,やはり sort -n というところか.

GNU shuf が mt-shuf と同じように動くのに気がついたので,実験結果に追加しておいた.GNU shuf と mt-shuf の速度差は,恐らくコンパイラの最適化の加減や例外処理等のオーバーヘッドの差と考えられる.