Suffix Trieを使って文字列マッチングする

文字列マッチングを行うためのアルゴリズムとして、Suffix Trieを使った探索というものがある。これはテキストからSuffix Trieという構造を作り、パターンをつかってそれを辿ることで、パターンの長さmに対して、O(m)の計算量で探索できるものである。

今回はJavaでSuffix Trieを使った探索をしてみた。

トライ木とパトリシア

先にトライ木とパトリシアについて紹介。

今回Suffix Trieという構造を調べていると、同じような構造としてSuffix Treeというものが出てきて混乱した。よく調べてみると、Suffixの集合をトライ木という構造で実装したものがSuffix Trieで、パトリシアという構造で実装したものがSuffix Treeらしい。

トライ木はnodeを連結していって、その枝に1文字を割り当てて辿れるようにした構造。トライ木は枝に1文字しか割り当てない構造上非常にメモリを食うので、複数文字格納するようにしてメモリ効率をよくしたものをパトリシアと言うらしい。詳しくは以下を参照。

枝に複数文字入れるということは、普通に実装が大変そうな予感がしたので、今回はトライ木構造を使うSuffix Trieの方を実装することにした。

Suffix Trieを使ったマッチングとは

Coursera の Algorithms on Strings 受けました - たにしきんぐダムのSuffix Trieの説明が分かりやすい。

まず、テキストの全てのSuffix集合を使ってトライ木を作る。これで例えば「nana」というテキストなら、上記の記事のような構造が出来る。

f:id:tanishiking24:20161218022331p:plain — 「Coursera の Algorithms on Strings 受けました」から引用

あとは「na」というパターンでマッチしようとすると、このトライ木を辿ったら以下のようになり、辿ったノードの下の全ての葉ノードの数字がマッチした位置となる。

f:id:shiba_yu36:20170108220909j:plain

他にも次の記事が参考になる。

実装

次のフローで文字列マッチングが出来ることが分かった。

与えられたテキストからトライ木を構築する
与えられたパターンでトライ木を辿り、辿った全ての子ノードを返す

というわけで実装は以下のとおりとなる。new SuffixTrie(text)とするとトライ木が構築され、そのインスタンスでsearchPatternするとpositionのリストが返ってくるようになっている。 https://github.com/shibayu36/algorithm-study/tree/master/algorithms-on-string/src/main/java あたりにコードは置いてある。

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

// new SuffixTrie(text).searchPattern(pattern)とすれば
// text中からpatternがどこに出現するかpositionのリストを返してくれる
public class SuffixTrie {
  public String text;
  public SuffixTrieNode root;

  public SuffixTrie(String text) {
    // 番兵を入れておけば、葉ノードからさらに枝が生えることがなくなるので
    // 構築や探索が楽になる。
    this.text = text + "\0"; // マッチングしやすいように番兵を入れておく
    this.root = new SuffixTrieNode();

    // 全てのSuffixを一つずつTrie木に追加していく
    for (int i = 0; i < this.text.length(); i++) {
      this.insertSuffix(new Suffix(this.text, i));
    }
  }

  // Suffix一つをTrie木に追加する
  private void insertSuffix(Suffix suffix) {
    SuffixTrieNode node = this.root;

    // 1文字ずつrootから辿りながらnodeを作成していく
    for (int i = 0; i < suffix.length(); i++) {
      Character c = suffix.charAt(i);
      if (node.children.containsKey(c)) {
        // 既に辿る枝があれば辿る
        node = node.children.get(c);
      }
      else {
        // なければ次のノードを作る
        SuffixTrieNode newNode = new SuffixTrieNode();
        node.children.put(c, newNode);
        node = newNode;
      }
    }

    // 辿った最後が葉なので、そこにSuffixのpositionを入れておく
    node.position = suffix.index;
  }

  // patternから出現するpositionのリストを返す
  public List<Integer> searchPattern(String pattern) {
    // まずpatternからマッチするnodeを取得する
    SuffixTrieNode matched = this.searchNode(pattern);
    if (matched != null) {
      // 取得したnodeの全ての葉ノードを取得すれば、positionのリストが得られる
      return matched.getAllLeafNodes().stream().map(
          node -> node.position
      ).collect(Collectors.toList());
    }
    else {
      return new ArrayList<>();
    }
  }

  // patternを辿ったnodeを返す
  // nodeが見つからなければnull
  private SuffixTrieNode searchNode(String pattern) {
    SuffixTrieNode node = this.root;
    for (int i = 0; i < pattern.length(); i++) {
      Character c = pattern.charAt(i);
      if (!node.children.containsKey(c)) {
        return null;
      }
      else {
        node = node.children.get(c);
      }
    }
    return node;
  }
}

class SuffixTrieNode {
  public int position; // どこからのsuffixの葉なのか
  public Map<Character, SuffixTrieNode> children = new HashMap<>();

  public SuffixTrieNode() {}

  public boolean isLeaf() {
    return this.children.isEmpty();
  }

  // あるノードの全ての葉ノードを返す
  public List<SuffixTrieNode> getAllLeafNodes() {
    if (this.isLeaf()) {
      List<SuffixTrieNode> nodes = new ArrayList<>();
      nodes.add(this);
      return nodes;
    }
    else {
      List<SuffixTrieNode> nodes = new ArrayList<>();
      for (SuffixTrieNode child : this.children.values()) {
        nodes.addAll(child.getAllLeafNodes());
      }
      return nodes;
    }
  }
}

class Suffix implements Comparable<Suffix> {
  public String text;
  public int index;

  public Suffix(String text, int index) {
    this.text = text;
    this.index = index;
  }
  public int length() {
    return text.length() - index;
  }
  public char charAt(int i) {
    return text.charAt(index + i);
  }

  public int compareTo(Suffix that) {
    if (this == that) return 0;  // optimization
    int n = Math.min(this.length(), that.length());
    for (int i = 0; i < n; i++) {
      if (this.charAt(i) < that.charAt(i)) return -1;
      if (this.charAt(i) > that.charAt(i)) return +1;
    }
    return this.length() - that.length();
  }
}

実際に動かしてみる

実装は出来たので実際に動かしてみる。

まず自分のブログ記事を https://github.com/shibayu36/algorithm-study/blob/master/algorithms-on-string/blog1.txt においた。これから検索をしてみる。試したコードは以下のとおり。

import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.Comparator;
import java.util.List;

public class SuffixTrieSearchSample {

  public static void main(String[] args) throws Exception {
    String text = "";

    try {
      byte[] fileContentBytes = Files.readAllBytes(Paths.get("./blog1.txt"));
      text = new String(fileContentBytes, StandardCharsets.UTF_8);
    } catch (Exception e) {}

    searchAndPrint(text, "基礎技術");
    System.out.println("");
    searchAndPrint(text, "モチベーション");
  }

  private static void searchAndPrint(String text, String pattern) {
    List<Integer> positions = new SuffixTrie(text).searchPattern(pattern);
    positions.sort(Comparator.naturalOrder());
    for (int position : positions) {
      int begin = Math.max(0, position - 5);
      int end = Math.min(text.length(), position + pattern.length() + 5);
      System.out.println(text.substring(begin, end));
    }
  }
}

結果は以下のようになり、正しく文字列マッチングできた。

基礎技術の学習のモ
寿命の長い基礎技術を、モチベ
の中でその基礎技術を実際のサ
開発の中で基礎技術との繋がり
寿命の長い基礎技術でも、モチ
ンを保って基礎技術を学習をす

術の学習のモチベーションをどう保つ
いないと、モチベーションをずっと保
が難しい
モチベーションが保てない
とりあえずモチベーションは保ちやす
きなければモチベーションを保つこと
礎技術を、モチベーションを保ちつづ
を得られ、モチベーションが増加する
で、さらにモチベーションを高められ
技術でも、モチベーションを保って学
どうやってモチベーションを保って基