2012/03/15

JavaでPatriciaTrieを実装してみた

書籍「日本語入力を支える技術」を読んでテンション上がったので、Javaでパトリシアトライを実装してみました(本当はビットベクタを実装したかったけどとりあえずダブル配列からやろうと思ってそれでも挫折してPatriciaTrieから入門した次第)。また、このエントリで紹介するコードはパトリシアトライの説明用で、全然最適化が行われていないので、メモリ効率も速度も良くないです。そのあたりを最適化したコードは、別のエントリで紹介する予定です。
(↓※アソシエイトIDあり)
ちなみにPatriciaはPractical Algorithm to Retrieve Information Coded in Alphanumericの略らしい(参考: wikipedia:基数木)。構造は至ってシンプルで、いろいろと実装例があります。
wikipediaのパトリシアトライの説明から図を借りると、その構造はこんな感じ:

ちょっとわかりにくいので、形を変えてみると:
空間効率的には、上図右側の各文字列の灰色で示した部分が他の文字列と共有されるので、かなりデータ量は小さくて済みます(wikipedia日本語タイトルデータで50%程度の効率化)。また、検索効率は、検索したい文字列を先頭から順に探して行くだけなので、最悪O(k)で済みます(kは文字列長)。実際は子ノードの探索方法によっては検索速度に影響が出るけど、それはまた後で。

構造がわかったところで、それを表現するクラスを定義します。

public class Node{
 // コンストラクタ、getter、setterは割愛
 private String letters;
 private List<node> children;
 private boolean terminated;
}

各ノードで保持する文字列は変数lettersに、子ノードは変数childrenに格納します。加えて、そのノードで文字列が終わっているかどうかを保持する変数terminatedを設けます。これは、例えば上図の文字列に加え、"rom"や"rubic"が加わった場合、ツリーの途中で文字列が終わることになるので、それに対応するため。上図に、文字列の終端という概念を加えて、あとツリーの途中で文字列が終わるケースを示すために"rom"と"rubic"も挿入されたとして、終端のノード(terminated==true)を太枠で表すと、下図のようになります。


終端については、終端文字列を設けて('$'とか)、それを終端となるノードにぶら下げるのが一般的のようなんですが、コードをシンプルにしたかったので、上記実装にしました(実際はいくつか工夫の余地はあり、今できるだけメモリを使わない方法をトライ中なんですが、それは完成した後に詳しく書く予定)。

さて、まずはある文字列がトライに含まれるかどうかを判定する、containsメソッドを実装してみます。考え方としては、
  1. 自ノードに文字列が設定されている場合、検索文字列との比較を行う。
  2. 自ノードの文字列だけでは比較が終わらない場合、子ノードを探す。
  3. 子ノードが見つかった場合、子ノードに対して同じ操作を行う
という感じで、検索対象文字列の先頭から、ノード内の文字列との比較を行い、どんどん子ノードに降りて行きます。コードは下記。

boolean contains(String word){
 int rest = word.length();
 int n = letters.length();
 if(n > rest) return false; // 自ノードの文字列のほうが長ければ比較失敗。
 for(int i = 0; i < n; i++){
  if(letters.charAt(i) != word.charAt(i)) return false;
 }
 if(n == rest) return terminated; // 文字列が同じであれば、終端かどうかを返す。
 word = word.substring(n); // 比較済み文字列を捨てる。
 if(children != null){
  char c = word.charAt(0); // 子ノードを探す。(*1)
  for(Node n : children){
   if(n.letters.charAt(0) == c){
    return n.contains(word); // 見つかれば、containsを呼び出す。
   }
  }
 }
 return false;
}
自ノードの文字列(letters)、子ノード(children)の順で処理を行なっています。子ノードの検索(*1)は線形探索していて、当然遅いので、実用で使うには、頻度順に並べるとか、childrenはソートしておいて2分探索を使うとか、ハッシュを使うとかする必要がありますが、今は割愛。データ量が膨大になる場合は、各ノードでハッシュを持つとメモリ消費が問題になる可能性があるので、2分探索か頻度順が現実的かと思います(そのぶん構築コストは高いけど)。実際wikipediaの日本語タイトルを挿入してみると、線形探索と2分探索では、処理にかかる時間の桁が違いました(50万件の検索で線形探索1762ms, 2分探索201ms)。

ちなみに2分探索だと、コードは下記のようになります。
boolean contains(String word){
 int rest = word.length();
 int n = letters.length();
 if(n > rest) return false; // 自ノードの文字列のほうが長ければ比較失敗。
 for(int i = 0; i < n; i++){
  if(letters.charAt(i) != word.charAt(i)) return false;
 }
 if(n == rest) return terminated; // 文字列が同じであれば、終端かどうかを返す。
 word = word.substring(n); // 比較済み文字列を捨てる。
 if(children != null){
  char c = word.charAt(0); // 2分探索で子ノードを探す
  int end = children.size();
  int start = 0;
  while(start < end){
   int i = (start + end) / 2;
   Node child = children.get(i);
   int d = c - child.getLetters().charAt(0);
   if(d == 0){
    return child.contains(word);
   }
   if(d < 0){
    end = i;
   } else if(start == i){
    break;
   } else{
    start = i;
   }
  }
 }
 return false;
}

次に要素の挿入です。おおまかに以下の流れで実装しました。
  1. ノード内の文字列(thisLetters)か比較対象の文字列(letters)の短い方の文字数分比較
  2. 全て比較成功
    1. thisLettersとlettersの長さが同じであれば、ノードを終端にして(terminated = true)終了(a)
    2. thisLettersの方が長ければ、ノードを分割して終了(b)
    3. lettersの方が長ければ、比較済み文字列を捨てる
    4. 新しいlettersの先頭文字とマッチする先頭文字を持つ子ノードを探す
    5. 存在すれば、その子ノードに対してinsertChildを呼び出して終了(c)
    6. 存在しなければ、新しい子ノードを作成して追加し終了(d)
  3. 途中で比較が止まった場合、その位置でノードを分割し、lettersの残りも子ノードとして追加する。(e)
コードは下記のようになります。
public void insertChild(String letters){
   int n = Math.min(letters.length(), this.letters.length());
   int i = 0;
   while(i < n && (letters.charAt(i) == this.letters.charAt(i))) i++;
   if(i == n){
    if(letters.length() == this.letters.length()){
     terminated = true; // (a)
     return;
    }
    if(letters.length() < this.letters.length()){
     Node node = new Node( // (b)
       this.letters.substring(letters.length())
       , this.children, this.terminated);
     this.letters = letters;
     this.children = new ArrayList<Node>(1);
     this.children.add(node);
     ((ArrayList<Node>)this.children).trimToSize();
     this.terminated = true;
     return;
    }
    letters = letters.substring(i);
    int index = 0;
    for(Node child : children){
     int c = letters.charAt(0) - child.letters.charAt(0);
     if(c < 0) break;
     if(c == 0){
      child.insertChild(letters); // (c)
      return;
     }
     index++;
    }
    this.children.add(index, new Node(letters)); // (d)
    ((ArrayList<Node>)this.children).trimToSize();
    return;
   }
   String newLetter1 = this.letters.substring(0, i); // (e)
   String newLetter2 = this.letters.substring(i);
   String newLetter3 = letters.substring(i);
   List<Node> newChildren = new ArrayList<Node>(2);
   if(newLetter2.charAt(0) < newLetter3.charAt(0)){
    newChildren.add(new Node(newLetter2, this.children, this.terminated));
    newChildren.add(new Node(newLetter3));
   } else{
    newChildren.add(new Node(newLetter3));
    newChildren.add(new Node(newLetter2, this.children, this.terminated));
   }
   this.letters = newLetter1;
   this.terminated = false;
   this.children = newChildren;
  }

なお、子ノードは全て昇順でソートされるようにして、探索は線型探索を使っています(ここも2分探索にすれば高速化できる)。また、メモリ利用量を抑えるため、子ノードの配列は正味のサイズだけ確保するようにしています(その分メモリリアロケーションが増えて速度は落ちているはず)。

ちなみに、上記のコードでは、VMの最大メモリ容量が128MBだと、Wikipedia日本語タイトル(127万件)はメモリに載りません(180MB程度必要です。ちなみにlettersをchar[]に、childrenをArrayListではなくNode[]に変えれば載ります)。ですが、まぁ一応機能はするので、次はTrieの特徴であるcommon prefix searchを実装してみます。

コメントを投稿