話題のwat-arrayを使ってBurrows-Wheeler変換(BWT)してみた

先日PFIの岡野原氏によってwat-arrayというライブラリが公開された。

wat-array : wavelet木を利用した高速配列処理ライブラリ : Preferred Research Blog

このライブラリは内部でウェーブレット木(wavelet tree)という簡潔データ構造(succinct data structure)を使っている。このため文字列に対するrank()やselect()などの操作が効率的にできるようになっている。
・・・といっても馴染みのない人にとっては何が嬉しいのかピンと来ないかもしれない。そこでBurrows-Wheeler変換(BWT, Burrows-Wheeler Transform)を例にとってwat-arrayの使いみちを説明してみる。
Burrows-Wheeler変換というのはテキストを同じ文字が並びやすいように変換したもので、通常ランレングス圧縮の前処理などで用いられる。またSuffixArrayとほぼ同じ情報を持っているため、FM-Indexのように全文検索用のデータ構造としても用いられている。


まずはBurrows-Wheeler変換の説明。まずは与えられたテキストを1文字ずつずらして並べる。便宜上末尾には任意の文字よりも小さいマーカ(ここでは#)を入れる。例えばmississippiというテキストの場合

mississippi#
ississippi#m
ssissippi#mi
sissippi#mis
issippi#miss
ssippi#missi
sippi#missis
ippi#mississ
ppi#mississi
pi#mississip
i#mississipp
#mississippi

となる。次にこれを辞書順ソートする。

#mississippi
i#mississipp
ippi#mississ
issippi#miss
ississippi#m
mississippi#
pi#mississip
ppi#mississi
sippi#missis
sissippi#mis
ssippi#missi
ssissippi#mi

このとき各行の最後の文字を並べた配列がBurrows-Wheeler変換されたテキストとなる。

ipssm#pissii

眺めてみるとなんとなく同じ文字が並んでいるのがわかる。Burrows-Wheeler変換ができたので元テキストの復元を考えたいところだが、ここで頭の整理を兼ねてBurrows-Wheeler変換の実装例を示す。あくまで例なので効率とかは考えていない。
実は変換のほうはwat-arrayは使わない。使うのは復元のほうであり今回の主題も復元のほう。なので変換はさくっと流す。

#include <vector>
#include <algorithm>
#include <cstring>
#include <stdint.h>
using namespase std;

struct less_char {
  bool operator() (const char *l, const char *r) {
    return strcmp(l, r) < 0;
  }
};

void text2bwt(const char *T, vector<uint64_t> &array) {
  const char *p = T;
  vector<const char *> v;
  while (*p != '\0') { v.push_back(p); p++; }
  sort(v.begin(), v.end(), less_char());

  array.clear();
  vector<const char *>::iterator i = v.begin();
  while (i != v.end()) {
    uint64_t pos = ((*i - T) + v.size() - 1) % v.size();
    array.push_back(T[pos]);
    i++;
  }
  return;
}

テキストTを受け取って、変換後のテキストをvector型で返す。文字列型で返さないのはwat-arrayが入力としてvectorを受け取るため。
次にテキストの復元を考える。以降、変換後のテキストを便宜上Tbwtとする。テキストの復元はTbwtの任意の要素に対して「元のテキスト上での直前の文字の位置」を返すような関数LF()というものをつくってあげれば良い。このLF()はwat-arrayの機能を使うことで非常に単純に実現できる。

#include "wat_array.hpp"
using namespace wat_array;

uint64_t LF(WatArray &wa, uint64_t i, uint64_t ch) {
  return wa.RankLessThan(ch, wa.length()) + wa.Rank(ch, i);
}

たったこれだけ。LF()の引数はwat-arrayのオブジェクトwaとTwbt上の位置i、それから位置iにある文字chの3つ。まずWatArray::RankLessThan()という関数を読んでいる。これは第2引数の位置より前方にchより(辞書順に)小さい文字が何回出たかを返すもの。今回は第2引数にはwa.length()しか使わない。つまりテキスト全体で何回出たか、だけを考える。例えば先程のmississippi#を変換したipssm#pissiiに対しては

wa.RankLessThan('#', wa.length()) => 0
wa.RankLessThan('i', wa.length()) => 1
wa.RankLessThan('m', wa.length()) => 5
wa.RankLessThan('p', wa.length()) => 6
wa.RankLessThan('s', wa.length()) => 8

となる。WatArray()::Rank()は第2引数の位置より前方にchが何回出たかを返すもの。ipssm#pissiiに対しては

wa.Rank('#', 3) => 0
wa.Rank('i', 3) => 1
wa.Rank('m', 3) => 0
wa.Rank('p', 3) => 1
wa.Rank('s', 3) => 1

wa.Rank('#', wa.length()) => 1
wa.Rank('i', wa.length()) => 4
wa.Rank('m', wa.length()) => 1
wa.Rank('p', wa.length()) => 2
wa.Rank('s', wa.length()) => 4

というふうになる。これで何故、元テキスト上での直前の位置がわかるのかを説明する。ここでTbwtをつくるときにテキストを1文字ずつずらして辞書順ソートしたことを思い出す。

#mississippi
i#mississipp
ippi#mississ
issippi#miss
ississippi#m
mississippi#
pi#mississip
ppi#mississi
sippi#missis
sissippi#mis
ssippi#missi
ssissippi#mi

この行列をWとする。例えば、8行目のppi#mississiに注目する。この行の末尾のiの直前に来る文字を探すには、ppi#mississiを右に1文字ずつずらしiを先頭にくっつけたippi#mississを探せば良い。Tbwtだけが手元にある状況でippi#mississを探すことを考える(手元にはipssm#pissiiという末尾の文字を並べた配列しかないことに注意。つまり明示的にippi#mississを探すぞ、といことはわからず、先頭がiのiXXXXXXXXXXXという文字列を探すという情報しかない)。
Wの各行は辞書順に並んでいるということを利用する。すると先頭の文字が同じ行は連続していて、かつ最初に先頭が文字chになる行は、chよりも小さい文字の総数+1行目になることがわかる。「chよりも小さい文字の総数」はwa.RankLessThan(ch, wa.length())を使えば求められる。ippi#mississは先頭の文字がiなので、wa.RankLessThan('i', wa.length())=1となるので2行目がiで始まる文字列の開始位置だとわかる(waは最初の行が0なのでRankLessThan()=1は2行目)。
Wのうちiで始まる文字列は4行ある。このうちどれが正しい行なのかを考える。そもそもiで始まる行を探しているということは、最初に注目した行の末尾がiで終わっているということ(今回は8行目のppi#mississiに注目していた)。つまり

#mississippi
ppi#mississi
ssippi#missi
ssissippi#mi

の4行。これらの文字列をそれぞれ右に1文字ずつずらしてiを先頭に付けると

i#mississipp
ippi#mississ
issippi#miss
ississippi#m

これはW上での並びと同じ。なぜ同じになるかというと、先頭の文字が全部iなので、辞書順での並びを決定するのは2文字目以降だから。このため先頭のiを末尾につけた文字列と並び順は同じになる。
さてWの各行について末尾の文字が同じ行は、末尾の文字を先頭につけた行と並び順が同じになるという性質を使う。この性質を使えば、注目していたWの8行目のppi#mississiという行が、末尾がiの行のうち何番目なのかがわかれば、iを先頭にうつした行が先頭の文字がiの行のうち何番目のものか、ということがわかる。
末尾がchの行jについて、それが末尾がchの行のうち幾つ目のものかを知るにはwa.Rank(ch, j)を使えば良い。今回の例ではwa.Rank('i', 7)=1となるので先頭がiの行は自分より前に1行あることがわかる(waは最初の行が0なのでRank()の第2引数は8ではなく7)。
以上を纏めると

wa.RankLessThan('i', wa.length()) + wa.Rank('i', 8) = 1 + 1 = 2

なのでWの3行目に求めるべきippi#mississがあることがわかる。この結果Tbwtの8番目の要素iの直前に来るのは3番目の要素sであることがわかる。
よってLF()関数で直前の文字が復元できることがわかったので末尾の文字#から順番にLF()を使って1文字ずつ復元していけば元テキストが得られることがわかる。つまり以下のようなコードが書ける。(中略)の箇所には、上で説明したtext2bwt()とLF()のコードをコピペすること。

#include <iostream>

(中略)

int main(int argc, char **argv) {
  vector<uint64_t> array;
  text2bwt(argv[1], array);
  
  vector<uint64_t>::iterator j = array.begin();
  while (j != array.end()) { cout << char(*j); j++; }
  cout << endl;

  WatArray wa;
  wa.Init(array);
  array.clear();
  uint64_t ch = '#';
  uint64_t i  = wa.Select(ch, 1);
  do {
    array.insert(array.begin(), ch);
    i = LF(wa, i, ch);
  } while ((ch = wa.Lookup(i)) != '#');

  j = array.begin();
  while (j != array.end()) { cout << char(*j); j++; }
  cout << endl;
  return 0;
}

これをtest.ccなどとしてコンパイルする。

$$ svn checkout http://wat-array.googlecode.com/svn/trunk wat-array
$$ g++ -Iwat-array/src test.cc wat-array/src/*.cpp -o test

$$ ./test mississippi#
ipssm#pissii
mississippi#

きちんとmississippi#が復元できていることがわかる。説明を書いたら結構長くなってしまったけど、wat-arrayを使うことでコードは非常に短く綺麗に書くことができたよ。というのが今回のポイント。