Trie木

できること

  • 文字列の追加。文字列の長さが Nの場合 \mathcal{O}(N)
  • 入力された文字列と同じ文字列の個数カウント。文字列の長さが Nの場合 \mathcal{O}(N)。これは何も考えずにsetとかに入れて二分探索すると文字列の個数 Mに対して \mathcal{O}(N \log M)かかるので、 Nと比べて Mがかなり大きい時に有利。
  • 入力された文字列をprefixとするような個数のカウント。文字列の長さが Nの場合 \mathcal{O}(N)。これはsetとかで実装するのは難しい。愚直に全探索すると、文字列の個数を Mとしたときに \mathcal{O}(NM)かかってしまう。

仕組み

1つの文字を1つのノードに対応させ、文字の連なりを連結リストとして表現する。prefixが共通な文字列は同じノードを再利用することでノード数を削減している。文字列の終わりのノードに印をつけておくことで入力文字列と一致する文字列があるかを判定できる。また、それぞれも文字はbase文字(例えば小文字アルファベットであればbase="a")からの引き算によって数字(以降文字数字と呼ぶ)として扱う(例えばcの場合は2に対応させる)。
図では文字列の終点に対応するノードのみにacceptをインクリメントしているが、途中のノードでもインクリメントすれば、完全一致ではなくprefix一致の個数も取得することができる。

f:id:salpik:20210117130357p:plain
Trie木の構造

実装

  • ノードのデータ構造
    以上の機能を実現するため、各ノードは以下のような構造体として扱うのが良い。
typedef struct TrieNode {
  int id;  // 各ノードが持つユニークID。生成された順番でインクリメントIDが付与される
  std::vector<int>
      next;  // 各文字数字について、接続しているノードのIDをつめる配列。(非接続時は-1)
  int end;   // この文字で終了した文字列の個数。初期値は0。
  char base;  // 最初の文字(小文字アルファベットを扱うのであれば'a')
  int string_varierty_size;  // 想定する文字の種類数(小文字アルファベットを扱うのであれば26)
  // コンストラクタの定義、TrieNodeのデータを作る時にこの3つの引数は必ず指定しないといけない
  TrieNode(int id_, char base_, int string_varierty_size_)
      : id(id_), base(base_), string_varierty_size(string_varierty_size_) {
    next = std::vector<int>(string_varierty_size, -1);
    end = 0;
    middle = 0;
  };
} TrieNode;

例えば図のケースの場合、nodes[2]に着目すると、nodes[2].next[2] = 5、他(nodes[0]、nodes[1], nodes[3],,,)は-1である。

  • 文字列一致個数算出クエリ
    文字列の一致は、Trie木を辿っていって、停止したところのacceptの値を見れば良い。また、途中でTrie木が途切れている場合は返り値を0にして早期リターンすれば良い。

  • 文字列の追加クエリ
    Trie木に沿ってノードを辿っていき、新しくnodeを追加する必要がある場合はノードを生成していけばよい。最後のノードのaccept値をインクリメントする必要がある。

クラスの構成

class TrieTree {
 public:
  std::vector<TrieNode> nodes;
  TrieTree(char base_, int string_varierty_size_)
      : base(base_), string_varierty_size(string_varierty_size_) {
    // 最初に親ノードを1つ追加しておくことを忘れない
    struct TrieNode firstnode(0, base_, string_varierty_size_);
    nodes.push_back(firstnode);
  };
  // 入力文字列と一致する文字列の個数をカウント
  int CountMatchString(std::string s) {
    TrieNode lastnode = GetLastNode(s);
    return lastnode.end;  // GetLastNodeでノードが見つからない場合は0を返す
  }
  void AddString(std::string s) {
    int nownode = 0;  // 今いるノードのID
    int nowindex = 0;  // 文字列sの何番目の文字を今チェックしているか。
    while (nowindex < s.length()) {
      int mojinum = Convert2Num(s[nowindex]);
      if (nodes[nownode].next[mojinum] == -1) {
        // 対応するノードがないので新しくノードを追加する
        TrieNode newnode(nodes.size(), base, string_varierty_size);
        // 新しく生成したノードを既存のノードに接続する
        nodes[nownode].next[mojinum] = newnode.id;
        // ノードを新しく作ったら配列に登録することを忘れない
        nodes.push_back(newnode);
      }
      // ノードを遷移させる前に、通過した文字数カウンターをインクリメント
      nodes[nownode].middle++;
      // ノードが存在しない場合は新しいノードを接続させたので、現在のノードIDを通常通り進めることができる
      ;
      nownode = nodes[nownode].next[mojinum];
      nowindex++;
    }
    // 最後にノードのaccept要素をインクリメントすることを忘れない
    nodes[nownode].end++;
    return;
  }

 private:
  char base;
  int string_varierty_size;
  int Convert2Num(char c) {
    int ret = int(c) - int(base);
    if (ret < 0 && ret >= string_varierty_size) {
      std::cerr << "Error: input char (" << c << ") exceeds expected region ["
                << base << ", " << base + string_varierty_size << "]"
                << std::endl;
      return 0;
    }
    return ret;
  }
  // 文字列に対応する最後のノードを返す
  TrieNode GetLastNode(std::string s) {
    int nownode = 0;  // 今いるノードのID
    int nowindex = 0;  // 文字列sの何番目の文字を今チェックしているか。
    while (nowindex < s.length()) {
      nownode = nodes[nownode].next[Convert2Num(s[nowindex])];
      if (nownode == -1) {
        return TrieNode(-1, base, string_varierty_size);  // 見つからないケース
      }
      nowindex++;
    }
    return nodes[nownode];
  }
};

入力文字列をprefixにもつ文字個数の算出

上記では文字列の最後のノードにのみaccept値をインクリメントしていたが、新たな変数を用意し途中のノードでもその値をインクリメントすれば、その値はそれをprefixにもつ文字列の個数を表していることになる。すなわち以下のような関数を用意すればよい。

  // 入力文字列をprefixとする文字列の個数をカウント
  int CountPrefixString(std::string s) {
    TrieNode lastnode = GetLastNode(s);
    return lastnode.middle;  // GetLastNodeでノードが見つからない場合は0を返す
  }

参考

トライ木(Trie木) の解説と実装【接頭辞(prefix) を利用したデータ構造】 | アルゴリズムロジック