Trie
(string/trie.hpp)

• View this file on GitHub
• Last update: 2026-07-15 18:17:55+09:00
• Include: #include "string/trie.hpp"

文字列の集合を trie で管理するデータ構造． 各ノードには，そのノードを通る登録文字列数と，そのノードで終端する文字列の ID 列を持つ．

Trie<char_size, char_start> として使う． 扱える文字は char_start 以上 char_start + char_size 未満であることを要求する．

• add(s, id = -1)：文字列 s を追加し，終端ノードを返す．id = -1 のとき，追加前の登録文字列数を ID として使う．
• insert(s, id = -1)：s が存在しなければ追加して true を返す．既に存在する場合は何もせず false を返す．
• erase(s)：s が存在すれば 1 個削除して true を返す．存在しない場合は何もせず false を返す．
• erase(s, id)：s の ID 列に id が存在すればそれを削除して true を返す．存在しない場合は何もせず false を返す．
• size(), count()：登録されている文字列数を返す．
• node_size()：trie のノード数を返す．
• empty()：登録されている文字列がないか判定する．
• count(s)：s と一致する登録文字列数を返す．
• count_prefix(s)：s を prefix に持つ登録文字列数を返す．
• contains(s)：s が存在するか判定する．
• find_ids(s)：s と一致する登録文字列の ID 列を返す．
• prefix_match(s)：s の prefix として登録されている文字列の ID 列を返す．
• enumerate_prefix(s)：s を prefix に持つ登録文字列の ID 列を返す．
• longest_prefix(s)：s の prefix として登録されている文字列の最大長を返す．存在しない場合は -1．
• root_node()：根ノードの index を返す．
• find_node(s)：s に対応するノードの index を返す．存在しない場合は -1．
• next_node(v, c)：ノード v から文字 c で遷移した先の index を返す．存在しない場合は -1．
• node(v)：ノード v を返す．
• restore(v)：ノード v に対応する文字列を復元する．
• clear()：すべての文字列を削除する．

erase はノードを物理的には削除しない．

add, insert, erase, count, count_prefix, contains, longest_prefix, find_node, restore は文字列長を $L$ として $O(L)$ 時間． ただし erase(s, id) は s と一致する ID の個数にも線形時間． find_ids(s), prefix_match(s) は返す ID 数を $K$ として $O(L+K)$ 時間． enumerate_prefix(s) は走査する部分木のサイズと返す ID 数に対して線形時間．

Verified with

• verify/string/UNIT_trie.test.cpp

Code

#pragma once

template <int char_size = 26, char char_start = 'a'>
struct Trie {
  struct Node {
    int prev, count, c;
    int next[char_size];
    vector<int> ids;
    Node(int prev, int c) : prev(prev), count(0), c(c) {
      memset(next, -1, sizeof(next));
    }
  };

  Trie() : root(make_node()) {}

  int add(const string& str, int id = -1) {
    if (id == -1) id = size();
    int x = root;
    nodes[x].count++;
    for (auto c : str) {
      x = get_child(x, char_index(c));
      nodes[x].count++;
    }
    nodes[x].ids.emplace_back(id);
    return x;
  }
  bool insert(const string& str, int id = -1) {
    if (contains(str)) return false;
    add(str, id);
    return true;
  }
  bool erase(const string& str) {
    int x = find_node(str);
    if (x == -1 || nodes[x].ids.empty()) return false;
    nodes[x].ids.pop_back();
    subtract_path(x);
    return true;
  }
  bool erase(const string& str, int id) {
    int x = find_node(str);
    if (x == -1) return false;
    auto it = find(nodes[x].ids.begin(), nodes[x].ids.end(), id);
    if (it == nodes[x].ids.end()) return false;
    nodes[x].ids.erase(it);
    subtract_path(x);
    return true;
  }
  int size() const { return nodes[root].count; }
  int count() const { return size(); }
  int node_size() const { return nodes.size(); }
  bool empty() const { return size() == 0; }
  int count(const string& str) const {
    int x = find_node(str);
    return x == -1 ? 0 : nodes[x].ids.size();
  }
  int count_prefix(const string& str) const {
    int x = find_node(str);
    return x == -1 ? 0 : nodes[x].count;
  }
  bool contains(const string& str) const { return count(str) > 0; }
  vector<int> find_ids(const string& str) const {
    int x = find_node(str);
    return x == -1 ? vector<int>() : nodes[x].ids;
  }
  vector<int> prefix_match(const string& str) const {
    vector<int> ret;
    int x = root;
    append(ret, nodes[x].ids);
    for (auto c : str) {
      x = next_node(x, c);
      if (x == -1) break;
      append(ret, nodes[x].ids);
    }
    return ret;
  }
  vector<int> enumerate_prefix(const string& str) const {
    vector<int> ret;
    int x = find_node(str);
    if (x != -1) dfs_ids(x, ret);
    return ret;
  }
  int longest_prefix(const string& str) const {
    int x = root, ret = nodes[x].ids.empty() ? -1 : 0;
    for (int i = 0; i < (int)str.size(); i++) {
      x = next_node(x, str[i]);
      if (x == -1) break;
      if (!nodes[x].ids.empty()) ret = i + 1;
    }
    return ret;
  }
  int root_node() const { return root; }
  int find_node(const string& str) const {
    int x = root;
    for (auto c : str) {
      x = next_node(x, c);
      if (x == -1) return -1;
    }
    return x;
  }
  int next_node(int index, char c) const {
    return index == -1 ? -1 : nodes[index].next[char_index(c)];
  }
  const Node& node(int index) const { return nodes[index]; }
  string restore(int index) const {
    string ret;
    while (index != root) {
      ret.push_back(char_start + nodes[index].c);
      index = nodes[index].prev;
    }
    reverse(ret.begin(), ret.end());
    return ret;
  }
  void clear() {
    nodes.clear();
    root = make_node();
  }

 private:
  vector<Node> nodes;
  int root;
  int make_node(int prev = -1, int c = -1) {
    nodes.push_back(Node(prev, c));
    return nodes.size() - 1;
  }
  int get_child(int index, int char_index) {
    if (nodes[index].next[char_index] == -1) {
      int child = make_node(index, char_index);
      nodes[index].next[char_index] = child;
    }
    return nodes[index].next[char_index];
  }
  int char_index(char c) const {
    int x = c - char_start;
    assert(0 <= x && x < char_size);
    return x;
  }
  void subtract_path(int index) {
    while (index != -1) {
      nodes[index].count--;
      index = nodes[index].prev;
    }
  }
  void append(vector<int>& a, const vector<int>& b) const {
    a.insert(a.end(), b.begin(), b.end());
  }
  void dfs_ids(int index, vector<int>& ret) const {
    append(ret, nodes[index].ids);
    for (int c = 0; c < char_size; c++)
      if (nodes[index].next[c] != -1) dfs_ids(nodes[index].next[c], ret);
  }
};

/**
 * @brief Trie
 * @docs docs/string/trie.md
 */

#line 2 "string/trie.hpp"

template <int char_size = 26, char char_start = 'a'>
struct Trie {
  struct Node {
    int prev, count, c;
    int next[char_size];
    vector<int> ids;
    Node(int prev, int c) : prev(prev), count(0), c(c) {
      memset(next, -1, sizeof(next));
    }
  };

  Trie() : root(make_node()) {}

  int add(const string& str, int id = -1) {
    if (id == -1) id = size();
    int x = root;
    nodes[x].count++;
    for (auto c : str) {
      x = get_child(x, char_index(c));
      nodes[x].count++;
    }
    nodes[x].ids.emplace_back(id);
    return x;
  }
  bool insert(const string& str, int id = -1) {
    if (contains(str)) return false;
    add(str, id);
    return true;
  }
  bool erase(const string& str) {
    int x = find_node(str);
    if (x == -1 || nodes[x].ids.empty()) return false;
    nodes[x].ids.pop_back();
    subtract_path(x);
    return true;
  }
  bool erase(const string& str, int id) {
    int x = find_node(str);
    if (x == -1) return false;
    auto it = find(nodes[x].ids.begin(), nodes[x].ids.end(), id);
    if (it == nodes[x].ids.end()) return false;
    nodes[x].ids.erase(it);
    subtract_path(x);
    return true;
  }
  int size() const { return nodes[root].count; }
  int count() const { return size(); }
  int node_size() const { return nodes.size(); }
  bool empty() const { return size() == 0; }
  int count(const string& str) const {
    int x = find_node(str);
    return x == -1 ? 0 : nodes[x].ids.size();
  }
  int count_prefix(const string& str) const {
    int x = find_node(str);
    return x == -1 ? 0 : nodes[x].count;
  }
  bool contains(const string& str) const { return count(str) > 0; }
  vector<int> find_ids(const string& str) const {
    int x = find_node(str);
    return x == -1 ? vector<int>() : nodes[x].ids;
  }
  vector<int> prefix_match(const string& str) const {
    vector<int> ret;
    int x = root;
    append(ret, nodes[x].ids);
    for (auto c : str) {
      x = next_node(x, c);
      if (x == -1) break;
      append(ret, nodes[x].ids);
    }
    return ret;
  }
  vector<int> enumerate_prefix(const string& str) const {
    vector<int> ret;
    int x = find_node(str);
    if (x != -1) dfs_ids(x, ret);
    return ret;
  }
  int longest_prefix(const string& str) const {
    int x = root, ret = nodes[x].ids.empty() ? -1 : 0;
    for (int i = 0; i < (int)str.size(); i++) {
      x = next_node(x, str[i]);
      if (x == -1) break;
      if (!nodes[x].ids.empty()) ret = i + 1;
    }
    return ret;
  }
  int root_node() const { return root; }
  int find_node(const string& str) const {
    int x = root;
    for (auto c : str) {
      x = next_node(x, c);
      if (x == -1) return -1;
    }
    return x;
  }
  int next_node(int index, char c) const {
    return index == -1 ? -1 : nodes[index].next[char_index(c)];
  }
  const Node& node(int index) const { return nodes[index]; }
  string restore(int index) const {
    string ret;
    while (index != root) {
      ret.push_back(char_start + nodes[index].c);
      index = nodes[index].prev;
    }
    reverse(ret.begin(), ret.end());
    return ret;
  }
  void clear() {
    nodes.clear();
    root = make_node();
  }

 private:
  vector<Node> nodes;
  int root;
  int make_node(int prev = -1, int c = -1) {
    nodes.push_back(Node(prev, c));
    return nodes.size() - 1;
  }
  int get_child(int index, int char_index) {
    if (nodes[index].next[char_index] == -1) {
      int child = make_node(index, char_index);
      nodes[index].next[char_index] = child;
    }
    return nodes[index].next[char_index];
  }
  int char_index(char c) const {
    int x = c - char_start;
    assert(0 <= x && x < char_size);
    return x;
  }
  void subtract_path(int index) {
    while (index != -1) {
      nodes[index].count--;
      index = nodes[index].prev;
    }
  }
  void append(vector<int>& a, const vector<int>& b) const {
    a.insert(a.end(), b.begin(), b.end());
  }
  void dfs_ids(int index, vector<int>& ret) const {
    append(ret, nodes[index].ids);
    for (int c = 0; c < char_size; c++)
      if (nodes[index].next[c] != -1) dfs_ids(nodes[index].next[c], ret);
  }
};

/**
 * @brief Trie
 * @docs docs/string/trie.md
 */

Back to top page