JavaでPatriciaTrieで各種検索

前回からの続きで、今回は検索メソッドを実装してみます。実装するのはこの2つ。

• common prefix search
• predictive search

common prefix searchは、クエリ"東京国際フォーラム"が与えられた時に、"東"、"東京"、"東京国"、"東京国際フォーラム"をトライから列挙するという検索です(実際これらの単語はWikipedia日本語版にあります)。

上図で赤枠で示したノードをたどりながら、終端ノードを列挙していく処理になります。先頭からノードをたどり、ノードのラベルとクエリを比較し、ノードのラベルを比較し終えたら、終端ノードであれば列挙、そして子ノードから残りのクエリの比較を行う候補を探し、そのノードに対しても同じ処理をしていきます。一方向に辿っていくだけなので再帰は必要なく、単純なループで実現できます。コードを下記に示します。

public Iterable<String> commonPrefixSearch(String query) {
 List<String> ret = new ArrayList<String>(); // 結果配列
 int cur = 0; // query内の照合位置
 Node node = root;
 while(node != null){
  String letters = node.getLetters();
  if(letters.length() > (query.length() - cur)) return ret; // 残りの文字列よりノードのラベルのほうが長ければマッチしないので終了
  for(int i = 0; i < letters.length(); i++){
   if(letters.charAt(i) != query.charAt(cur + i)) return ret;
  }
  if(node.isTerminated()){
   ret.add(query.substring(0 , cur + letters.length()));
  }
  cur += letters.length();
  if(query.length() == cur) return ret;
  node = node.getChild(query.charAt(cur));
 }
 return ret;
}

単純なwhileループで、着目するノード(変数node)を書き換えながら処理を行なっています。
predictive searchは、指定された文字列(prefix)を前方一致検索します。例えば"東京国"が指定されれば、"東京国"、"東京国税局"、"東京国際フォーラム"、"東京国際マラソン"が列挙されます。

上図赤枠のノードが列挙対象になります。まず指定されたprefixを含むノードを探し、見つかったノード以降のノードを全て列挙します。この列挙は部分木全部が対象になるので、再帰で実装するのが楽です(再帰が嫌な場合はQueue使ってもOK)。コードはこんな感じ。

/**
 * 部分木の列挙をする再帰関数。
 */
private static void enumLetters(Node node, String prefix, List<String> letters){
 List<node> children = node.getChildren();
 if(children == null) return;
 for(Node child : children){
  String text = prefix + child.getLetters();
  if(child.isTerminated()) letters.add(text);
  enumLetters(child, text, letters);
 }
}

public Iterable<String> predictiveSearch(String prefix) {
 char[] prefixChars = prefix.toCharArray();
 int cur = 0; // prefixChars内の照合位置
 Node node = root;
 while(node != null){
  String letters = node.getLetters();
  int n = Math.min(letters.length(), prefixChars.length - cur); // 共通部分の比較のためnは短い方に合わせる。
  for(int i = 0; i < n; i++){
   if(letters.charAt(i) != prefixChars[cur + i]){
    return Collections.emptyList();
   }
  }
  cur += n;
  if(prefixChars.length == cur){ // prefixCharsの比較が終われば、あとは部分木の列挙
   List<String> ret = new ArrayList<String>();
   prefix += letters.substring(n);
   if(node.isTerminated()) ret.add(prefix);
   enumLetters(node, prefix, ret);
   return ret;
  }
  // 比較がまだ終わって無ければ、子ノードを探して検索を継続
  node = node.getChild(prefixChars[cur]);
 }
 return Collections.emptyList();
}