marisa-trie における rank/select の実装

概要

rank/select は簡潔データ構造（Succinct Data Structures）の核になる関数です．ビット列の k ビット目までに含まれる 0, 1 の数を求めるのが rank，k 番目の 0, 1 の位置（Index）を求めるのが select であり，ビット列の密度（1 の割合）によって，いろいろな実装があります．

marisa-trie では，0, 1 の割合が極端に偏らないビット列を想定するとともに，32-bit 環境における性能の劣化を防ぐために，64-bit 整数を使わないようにしました．そのため，ほとんどの部分は以前に開発したライブラリからの流用ですが，新しく書き直した部分もあります．ちなみに，索引のサイズはビット列の長さ n bits に対して (1/4)n bits です．

基本

ビット列の実装

ビット列の格納には 32-bit 整数の配列を使っています．8-bit 整数や 16-bit 整数を使わない理由は，32bits をまとめて計算することで効率が良くなるからです．一方，64-bit 整数を使わない理由は，前述のとおり，32-bit 環境における性能の劣化を防ぐためです．

PopCount

marisa-trie において，PopCount は整数に含まれる 1 の数を求めるクラスです．以前にも説明したように（Succinct なトライの実験その 6 - やた＠はてな日記），整数全体に含まれる 1 の数だけではなく，下位 8bits, 16bits, 24bits に含まれる 1 の数もまとめて計算しておき，後で利用するというのがポイントです．よく知られている実装では捨ててしまいますが，「それを捨てるなんてとんでもない」ということもあるのです．

class PopCount {
 public:
  PopCount(UInt32 x) : value_() {
    x = (x & static_cast<UInt32>(0x55555555UL))
        + ((x >> 1) & static_cast<UInt32>(0x55555555UL));
    x = (x & static_cast<UInt32>(0x33333333UL))
        + ((x >> 2) & static_cast<UInt32>(0x33333333UL));
    x = (x + (x >> 4)) & static_cast<UInt32>(0x0F0F0F0FUL);

    // 下位 8bits, 16bits, 24bits に含まれる 1 の数を残したいので，
    // 右じゃなくて左にシフトします．
    x += x << 8;
    x += x << 16;
    value_ = x;
  }

  UInt32 lo8() const {
    return value_ & 0xFFU;
  }
  UInt32 lo16() const {
    return (value_ >> 8) & 0xFFU;
  }
  UInt32 lo24() const {
    return (value_ >> 16) & 0xFFU;
  }
  UInt32 lo32() const {
    return value_ >> 24;
  }

 private:
  UInt32 value_;
};

アセンブラは触らぬ神に祟りなしということで…．

rank

rank の索引

rank については，一定サイズのブロックごとに索引を持たせるという基本的な実装になっています．marisa-trie で採用しているのは，512bits のブロック毎に絶対値を保存し，64bits 毎にブロック内での相対値を保存するという方式です．ブロック単位の索引として用意したクラスが marisa::Rank です．

// lib/marisa/rank.h
class Rank {
 public:
  UInt32 abs() const;   // rank1(512k)
  UInt32 rel1() const;  // rank1(512k + 64) - rank1(512k)
  UInt32 rel2() const;  // rank1(512k + 128) - rank1(512k)
  UInt32 rel3() const;  // rank1(512k + 192) - rank1(512k)
  UInt32 rel4() const;  // rank1(512k + 256) - rank1(512k)
  UInt32 rel5() const;  // rank1(512k + 320) - rank1(512k)
  UInt32 rel6() const;  // rank1(512k + 384) - rank1(512k)
  UInt32 rel7() const;  // rank1(512k + 448) - rank1(512k)

 private:
  UInt32 abs_;
  UInt32 rel_lo_;
  UInt32 rel_hi_;
};

コメントに書いてあるように，abs がブロック先頭における絶対値，relX がブロック内での相対値を返すようになっています．メンバ変数の割り当ては，abs に abs_，rel1 〜 rel4 に rel_lo_，rel5 〜 rel7 に rel_hi_ となっています．relX の詳細は以下のとおりです．

メンバ関数	メンバ変数	マスク
rel1	rel_lo_	0x0000007F
rel2	rel_lo_	0x00007F80
rel3	rel_lo_	0x007F8000
rel4	rel_lo_	0xFF800000
rel5	rel_hi_	0x000001FF
rel6	rel_hi_	0x0003FE00
rel7	rel_hi_	0x07FC0000

64-bit 整数が使えると 9bits ずつの割り当てですっきりするのですが，32-bit 整数を使うために，少し面倒なことになっています．後，rel_hi_ の上から 5bits は使っていないので，少し無駄があります．

索引のサイズは，512bits 毎に 96bits なので，ビット列の長さを n bits とおくと，約 (3/16)n bits になります．

rank の実装

索引が用意できれば，後は rank1（1 のビットを数える関数）を実装するのみです．引数の値から対応するブロックを算出して，ブロック内での相対値を取り出し，残りを数えるという手順になります．

// lib/marisa/bitvector.cc
UInt32 BitVector::rank1(UInt32 i) const {
  // 対応するブロック（i / 512）を求めます．
  const Rank &rank = ranks_[i / 512];

  // ブロックの先頭における絶対値を取り出します．
  UInt32 offset = rank.abs();

  // ブロック内での相対値を取り出して，直前に取り出した絶対値に加えます．
  switch ((i / 64) % 8) {
    case 1: {
      offset += rank.rel1();
      break;
    }
    case 2: {
      offset += rank.rel2();
      break;
    }
    case 3: {
      offset += rank.rel3();
      break;
    }
    case 4: {
      offset += rank.rel4();
      break;
    }
    case 5: {
      offset += rank.rel5();
      break;
    }
    case 6: {
      offset += rank.rel6();
      break;
    }
    case 7: {
      offset += rank.rel7();
      break;
    }
  }

  // PopCount を使って残りを数えます．
  switch ((i / 32) % 2) {
    case 1: {  // 残りが 32bits 以上の場合，最初の 32bits を調べます．
      offset += PopCount(blocks_[(i / 32) - 1]).lo32();
    }
    case 0: {  // 残る (i % 32)bits を調べます．
      offset += PopCount(blocks_[i / 32]
          & ((static_cast<UInt32>(1) << (i % 32)) - 1)).lo32();
      break;
    }
  }
  return offset;
}

rank1 を実装できれば，rank0 は rank1 を使って簡単に定義できます．rank0 と rank1 で索引を個別に用意する必要はありません．

// bitvector.h
UInt32 rank0(UInt32 i) const {
  return i - rank1(i);
}

select

select の索引

select も同じように実装できればよいのですが，k 番目の 1 と (k+1) 番目の 1 の間に 0 がどれだけあるのかが分からないため，そうもいきません．~~聖闘士に同じ技は通用しないということです．~~

そこで，rank の索引を使って二分探索という，基本的な実装を踏襲することにしました．ただし，O(logn) という計算量は存外に大きいので，select1(512k), select0(512k) を保存しておき，探索範囲の絞り込みに利用しています．

索引のサイズは，1 が 512bits 毎に 32bits，0 が 512bits 毎に 32bits なので，ビット列の長さ n bits に対して，約 (1/16)n bits になります．rank の索引と併せると，約 (1/4)n bits です．

select1 と select0 は同じように実装できるため，以下では select1 の実装について説明します．

select の実装（512bits 単位の絞り込み）

select1（(i+1) 番目の 1 を見つける関数）の第一段階は，ブロック（512bits 単位）の絞り込みです．まず，select の索引を使ってブロックの範囲を限定し，rank の索引を使ってブロックを絞り込みます．

// lib/marisa/bitvector.cc
UInt32 BitVector::select1(UInt32 i) const {
  // i が 512 で割りきれるときは，select の索引として持っている値を
  // そのまま返します．select1(0) を高速化できる上に境界を潰せるので
  // 一石二鳥だと考えると得した気分になるかもしれません．
  UInt32 select_id = i / 512;
  if ((i % 512) == 0) {
    return select1s_[select_id];
  }

  // select の索引を使って，二分探索の範囲を絞り込みます．
  UInt32 begin = select1s_[select_id] / 512;
  UInt32 end = (select1s_[select_id + 1] + 511) / 512;

  // (i+1) 番目の 1 が出現するブロックを二分探索により求めます．
  // 求める位置がブロックの先頭になる場合については，
  // 既に解決されているので，気にする必要はありません．
  while (begin + 1 < end) {
    UInt32 middle = (begin + end) / 2;
    if (i < ranks_[middle].abs()) {
      end = middle;
    } else {
      begin = middle;
    }
  }
  UInt32 rank_id = begin;
  i -= ranks_[rank_id].abs();
  ...
}

select の実装（64bits 単位の絞り込み）

ブロック単位の絞り込みの後は，64bits 単位の絞り込みです．rank の索引を利用して，再び二分探索をおこないます．

// lib/marisa/bitvector.cc
UInt32 BitVector::select1(UInt32 i) const {
  ...
  // 書くのが面倒なだけで，やっていることは二分探索です．
  const Rank &rank = ranks_[rank_id];
  UInt32 block_id = rank_id * 16;
  if (i < rank.rel4()) {
    if (i < rank.rel2()) {
      if (i >= rank.rel1()) {
        block_id += 2;
        i -= rank.rel1();
      }
    } else if (i < rank.rel3()) {
      block_id += 4;
      i -= rank.rel2();
    } else {
      block_id += 6;
      i -= rank.rel3();
    }
  } else if (i < rank.rel6()) {
    if (i < rank.rel5()) {
      block_id += 8;
      i -= rank.rel4();
    } else {
      block_id += 10;
      i -= rank.rel5();
    }
  } else if (i < rank.rel7()) {
    block_id += 12;
    i -= rank.rel6();
  } else {
    block_id += 14;
    i -= rank.rel7();
  }
  ...
}

select の実装（仕上げ）

64bits 単位まで絞り込めたら，後は PopCount とテーブル（Succinct なトライの実験その 8（select() の計算にテーブルを導入） - やた＠はてな日記）を使います．まず，前半 32bits と後半 32bits のどちらに含まれるかを検証し，最終的に残った 32bits について，二分探索により 8bits 単位まで絞り込みます．最後は，あらかじめ作成しておいたテーブルを参照して終了となります．

// lib/marisa/bitvector.cc
UInt32 BitVector::select1(UInt32 i) const {
  ...
  UInt32 block = blocks_[block_id];
  PopCount count(block);

  // 前半の 32bits に含まれていないときは，後半の 32bits を調べます．
  if (i >= count.lo32()) {
    ++block_id;
    i -= count.lo32();
    block = blocks_[block_id];
    count = PopCount(block);
  }

  // PopCount により求めた下位 8bits, 16bits, 24bits に含まれる 1 の数を
  // を使って二分探索をおこないます．
  UInt32 bit_id = block_id * 32;
  if (i < count.lo16()) {
    if (i >= count.lo8()) {
      bit_id += 8;
      block >>= 8;
      i -= count.lo8();
    }
  } else if (i < count.lo24()) {
    bit_id += 16;
    block >>= 16;
    i -= count.lo16();
  } else {
    bit_id += 24;
    block >>= 24;
    i -= count.lo24();
  }

  // 後は 8x256 通りしかないので，あらかじめ計算しておいた結果を使います．
  return bit_id + SelectTable[i][block & 0xFF];
}

まとめ

地味な内容ばかりですけど，無視できないほどの影響があります．本格的に rank/select を使ってみようかなと思っている方のお役に立てれば幸いです．