library/sequence/Trie.hpp

• View this file on GitHub
• Last update: 2025-11-24 18:49:07+09:00
• Include: #include "library/sequence/Trie.hpp"

Depends on

• build/pch/stdc++.hpp
• library/algebra/group/Add.hpp
• library/sequence/ForString.hpp

Required by

• library/sequence/AhoCorasick.hpp

Verified with

• test/AOJ/2212.test.cpp
• test/yukicoder/430.test.cpp
• test/yukicoder/430_2.test.cpp

Code

#pragma once
#include "library/algebra/group/Add.hpp"
#include "library/sequence/ForString.hpp"
template <typename CHAR, int SIGMA, typename AbelMonoid = GroupAdd<int>>
class Trie {
  protected:
    using X = typename AbelMonoid::value_type;
    struct Node {
        std::array<int, SIGMA> nxt;
        int pre;
        X val, suffix_val; // suffix_val は自身を含まない
        Node(int pre)
            : pre(pre), val(AbelMonoid::unit()),
              suffix_val(AbelMonoid::unit()) {
            std::ranges::fill(nxt, -1);
        }
    };
    std::vector<Node> nodes;

  public:
    Trie() : nodes(1, Node(-1)) {}

    int &nxt(int now, const CHAR &a) { return nodes[now].nxt[a]; }
    const int &nxt(int now, const CHAR &a) const { return nodes[now].nxt[a]; }

    int add(const std::vector<CHAR> &v, X x = 1) {
        int now = 0;
        for (const CHAR &a : v) {
            assert(0 <= a and a < SIGMA);
            if (!~nxt(now, a)) {
                nxt(now, a) = nodes.size();
                nodes.emplace_back(now);
            }
            AbelMonoid::Rchop(nodes[now].suffix_val, x);
            now = nxt(now, a);
        }
        AbelMonoid::Rchop(nodes[now].val, x);
        return now;
    }
    int node_idx(const std::vector<CHAR> &v) const {
        // s の Node を返す　追加されて無ければ -1
        int now = 0;
        for (const CHAR &a : v) {
            if (!~nxt(now, a))
                return -1;
            now = nxt(now, a);
        }
        return now;
    }
    X val(const std::vector<CHAR> &v) {
        int id = node_idx(v);
        return (~id ? nodes[id].val : AbelMonoid::unit());
    }
    X &val(int node_id) { return nodes[node_id].val; }
    // vは含まない
    X prefix_prod(const std::vector<CHAR> &v) {
        int now = 0;
        X sum = AbelMonoid::unit();
        for (const CHAR &a : v) {
            if (!~nxt(now, a))
                break;
            AbelMonoid::Rchop(sum, nodes[now].val);
            now = nxt(now, a);
        }
        return sum;
    }
    // vは含まない
    X suffix_prod(const std::vector<CHAR> &v) const {
        int id = node_idx(v);
        return (~id ? nodes[id].suffix_val : AbelMonoid::unit());
    }
    std::vector<CHAR> restore(int node_id) {
        assert(0 <= node_id and node_id < nodes.size());
        std::vector<CHAR> res;
        while (~nodes[node_id].pre) {
            int pre = nodes[node_id].pre;
            for (int j = 0; j < SIGMA; j++)
                if (nxt(pre, j) == node_id) {
                    res.push_back(j);
                    break;
                }
            node_id = pre;
        }
        std::ranges::reverse(res);
        return res;
    }
    X prod() const { return nodes[0].suffix_val; }
    int size() const { return nodes.size(); }

    template <typename F>
    void query(const std::vector<CHAR> &v, const F &f, int l = 0, int r = -1) {
        if (r < 0)
            r = v.size();
        int now = 0;
        for (int i = l; i < r; i++) {
            now = nxt(now, v[i]);
            if (~now)
                f(now);
            else
                break;
        }
    }
};

#line 2 "library/algebra/group/Add.hpp"
template<typename X>
struct GroupAdd {
  using value_type = X;
  static constexpr X op(const X &x, const X &y) noexcept { return x + y; }
  static constexpr void Rchop(X&x, const X&y){ x+=y; }
  static constexpr void Lchop(const X&x, X&y){ y+=x; }
  static constexpr X inverse(const X &x) noexcept { return -x; }
  static constexpr X power(const X &x, long long n) noexcept { return X(n) * x; }
  static constexpr X unit() { return X(0); }
  static constexpr bool commute = true;
};
#line 2 "library/sequence/ForString.hpp"
template <char MARGIN> struct ForString {
    static constexpr char change(char c) { return c - MARGIN; }
    static constexpr char restore(char a) { return a + MARGIN; }

    static std::vector<char> change(const std::string &s) {
        std::vector<char> v(s.size());
        for (int i = 0; i < s.size(); i++)
            v[i] = change(s[i]);
        return v;
    }
    static std::string restore(const std::vector<char> &v) {
        std::string s(v.size(), '#');
        for (int i = 0; i < v.size(); i++)
            s[i] = restore(v[i]);
        return s;
    }
};
struct FSAa {
    static constexpr char change(char c) {
        return c <= 'Z' ? c - 'A' : 26 + c - 'a';
    }
    static constexpr char restore(char a) {
        return a < 26 ? 'A' : a - 26 + 'a';
    }
    static std::vector<char> change(const std::string &s) {
        std::vector<char> v(s.size());
        for (int i = 0; i < s.size(); i++)
            v[i] = change(s[i]);
        return v;
    }
    static std::string restore(const std::vector<char> &v) {
        std::string s(v.size(), '#');
        for (int i = 0; i < v.size(); i++)
            s[i] = restore(v[i]);
        return s;
    }
};
using FSA = ForString<'A'>;
using FSa = ForString<'a'>;
using FS0 = ForString<'0'>;

#ifdef STR
#define STRA(s)                                                                \
    STR(s##tomato);                                                            \
    auto s = FSA::change(s##tomato);
#define STRa(s)                                                                \
    STR(s##tomato);                                                            \
    auto s = FSa::change(s##tomato);
#define STR0(s)                                                                \
    STR(s##tomato);                                                            \
    auto s = FS0::change(s##tomato);
#endif
#line 4 "library/sequence/Trie.hpp"
template <typename CHAR, int SIGMA, typename AbelMonoid = GroupAdd<int>>
class Trie {
  protected:
    using X = typename AbelMonoid::value_type;
    struct Node {
        std::array<int, SIGMA> nxt;
        int pre;
        X val, suffix_val; // suffix_val は自身を含まない
        Node(int pre)
            : pre(pre), val(AbelMonoid::unit()),
              suffix_val(AbelMonoid::unit()) {
            std::ranges::fill(nxt, -1);
        }
    };
    std::vector<Node> nodes;

  public:
    Trie() : nodes(1, Node(-1)) {}

    int &nxt(int now, const CHAR &a) { return nodes[now].nxt[a]; }
    const int &nxt(int now, const CHAR &a) const { return nodes[now].nxt[a]; }

    int add(const std::vector<CHAR> &v, X x = 1) {
        int now = 0;
        for (const CHAR &a : v) {
            assert(0 <= a and a < SIGMA);
            if (!~nxt(now, a)) {
                nxt(now, a) = nodes.size();
                nodes.emplace_back(now);
            }
            AbelMonoid::Rchop(nodes[now].suffix_val, x);
            now = nxt(now, a);
        }
        AbelMonoid::Rchop(nodes[now].val, x);
        return now;
    }
    int node_idx(const std::vector<CHAR> &v) const {
        // s の Node を返す　追加されて無ければ -1
        int now = 0;
        for (const CHAR &a : v) {
            if (!~nxt(now, a))
                return -1;
            now = nxt(now, a);
        }
        return now;
    }
    X val(const std::vector<CHAR> &v) {
        int id = node_idx(v);
        return (~id ? nodes[id].val : AbelMonoid::unit());
    }
    X &val(int node_id) { return nodes[node_id].val; }
    // vは含まない
    X prefix_prod(const std::vector<CHAR> &v) {
        int now = 0;
        X sum = AbelMonoid::unit();
        for (const CHAR &a : v) {
            if (!~nxt(now, a))
                break;
            AbelMonoid::Rchop(sum, nodes[now].val);
            now = nxt(now, a);
        }
        return sum;
    }
    // vは含まない
    X suffix_prod(const std::vector<CHAR> &v) const {
        int id = node_idx(v);
        return (~id ? nodes[id].suffix_val : AbelMonoid::unit());
    }
    std::vector<CHAR> restore(int node_id) {
        assert(0 <= node_id and node_id < nodes.size());
        std::vector<CHAR> res;
        while (~nodes[node_id].pre) {
            int pre = nodes[node_id].pre;
            for (int j = 0; j < SIGMA; j++)
                if (nxt(pre, j) == node_id) {
                    res.push_back(j);
                    break;
                }
            node_id = pre;
        }
        std::ranges::reverse(res);
        return res;
    }
    X prod() const { return nodes[0].suffix_val; }
    int size() const { return nodes.size(); }

    template <typename F>
    void query(const std::vector<CHAR> &v, const F &f, int l = 0, int r = -1) {
        if (r < 0)
            r = v.size();
        int now = 0;
        for (int i = l; i < r; i++) {
            now = nxt(now, v[i]);
            if (~now)
                f(now);
            else
                break;
        }
    }
};

Back to top page