Trie
(src/string/trie.hpp)

• View this file on GitHub
• Last update: 2024-11-09 02:46:49+09:00
• Include: #include "src/string/trie.hpp"

Trie

集合に対する文字列の追加や検索を効率的に行うデータ構造です．

コンストラクタ

Trie<size_t X = 26, char margin = 'a'> trie(char c = '$', int p = -1)

• 文字の種類数が $X$ で最小の文字が $\mathrm{margin}$ であるような空の木 trie を作ります．

• c は根ノードの文字です．

• p は親ノードの番号です．

計算量

• $O(1)$

next

int& next(int i, int j)

基本的に内部用の関数です．

ノード $i$ からのびる $j$ 番目に小さい文字のノード番号を返します．
そのようなノードがない場合は $-1$ を返します．

制約

• $i$ に存在しないノードを指定しない．
• $0 \leq j < X$

計算量

• $O(1)$

insert

void insert(string s, int x)

trie に識別子 $x$ として文字列 $s$ を追加します．

計算量

$s$ の長さを $n$ として，

• $O(n)$

find

int find(string s)

文字列 $s$ の最後の文字に対応するノード番号を返します．
文字列 $s$ が trie に含まれていない場合は $-1$ を返します．

計算量

$s$ の長さを $n$ として，

• $O(n)$

move

int move(int pos, char c)

ノード $\mathrm{pos}$ からのびる文字 $c$ のノード番号を返します．

制約

• $\mathrm{pos}$ に存在しないノードを指定しない．

計算量

• $O(1)$

size

int size()

trie のノードの個数を返します．

計算量

• $O(1)$

idx

int idx(int pos)

ノード $\mathrm{pos}$ がある文字列の最後の文字に対応するノードである場合，その文字列の識別子を返します (2つ以上ある場合，最新のものを返します)．

制約

• $\mathrm{pos}$ に存在しないノードを指定しない．

計算量

• $O(1)$

idxs

vector<int> idxs(int pos)

ノード $\mathrm{pos}$ がある文字列の最後の文字に対応するノードである場合，その文字列の識別子をすべて返します．

制約

• $\mathrm{pos}$ に存在しないノードを指定しない．

計算量

返り値の配列の長さを $m$ として，

• $O(m)$

count

int count(int pos)

ノード $\mathrm{pos}$ 以下に格納されている文字列の個数を返します．

制約

• $\mathrm{pos}$ に存在しないノードを指定しない．

計算量

• $O(1)$

par

int par(int pos)

ノード $\mathrm{pos}$ の親ノードの番号を返します．

制約

• $\mathrm{pos}$ に存在しないノードを指定しない．

計算量

• $O(1)$

Depends on

• template (src/template/template.hpp)

Required by

• AhoCorasick (src/string/aho_corasick.hpp)

Verified with

• verify/yukicoder/430.test.cpp

Code

#pragma once
#include "../template/template.hpp"
template <size_t X = 26, char margin = 'a'>
struct Trie {
    struct Node {
        array<int, X> nxt;
        vector<int> idxs;
        int idx, count, parent;
        char key;
        Node(const char c, const int par)
            : idx(-1), count(0), parent(par), key(c) {
            fill(nxt.begin(), nxt.end(), -1);
        }
    };
    vector<Node> st;
    Trie(const char c = '$', const int p = -1) {
        st.emplace_back(c, p);
    }
    inline int& next(const int i, const int j) {
        assert(0 <= i and i < (int)st.size());
        assert(0 <= j and j < (int)X);
        return st[i].nxt[j];
    }
    void insert(const string& s, const int x) {
        int pos = 0;
        for(int i = 0; i < (int)s.size(); ++i) {
            ++st[pos].count;
            const int k = s[i] - margin;
            if(~next(pos, k)) {
                pos = next(pos, k);
                continue;
            }
            const int npos = st.size();
            next(pos, k) = npos;
            st.emplace_back(s[i], pos);
            pos = npos;
        }
        ++st[pos].count;
        st[pos].idx = x;
        st[pos].idxs.emplace_back(x);
    }
    int find(const string& s) {
        int pos = 0;
        for(int i = 0; i < (int)s.size(); ++i) {
            const int k = s[i] - margin;
            if(next(pos, k) < 0) return -1;
            pos = next(pos, k);
        }
        return pos;
    }
    int move(const int pos, const char c) {
        assert(0 <= pos and pos < (int)st.size());
        return next(pos, c - margin);
    }
    int size() const {
        return st.size();
    }
    int idx(const int pos) const {
        assert(0 <= pos and pos < (int)st.size());
        return st[pos].idx;
    }
    int count(const int pos) const {
        assert(0 <= pos and pos < (int)st.size());
        return st[pos].count;
    }
    int par(const int pos) const {
        assert(0 <= pos and pos < (int)st.size());
        return st[pos].parent;
    }
    vector<int> idxs(const int pos) const {
        assert(0 <= pos and pos < (int)st.size());
        return st[pos].idxs;
    }
};

#line 2 "src/template/template.hpp"
#include <bits/stdc++.h>
using namespace std;
using ll = long long;
using P = pair<long long, long long>;
#define rep(i, a, b) for(long long i = (a); i < (b); ++i)
#define rrep(i, a, b) for(long long i = (a); i >= (b); --i)
constexpr long long inf = 4e18;
struct SetupIO {
    SetupIO() {
        ios::sync_with_stdio(0);
        cin.tie(0);
        cout << fixed << setprecision(30);
    }
} setup_io;
#line 3 "src/string/trie.hpp"
template <size_t X = 26, char margin = 'a'>
struct Trie {
    struct Node {
        array<int, X> nxt;
        vector<int> idxs;
        int idx, count, parent;
        char key;
        Node(const char c, const int par)
            : idx(-1), count(0), parent(par), key(c) {
            fill(nxt.begin(), nxt.end(), -1);
        }
    };
    vector<Node> st;
    Trie(const char c = '$', const int p = -1) {
        st.emplace_back(c, p);
    }
    inline int& next(const int i, const int j) {
        assert(0 <= i and i < (int)st.size());
        assert(0 <= j and j < (int)X);
        return st[i].nxt[j];
    }
    void insert(const string& s, const int x) {
        int pos = 0;
        for(int i = 0; i < (int)s.size(); ++i) {
            ++st[pos].count;
            const int k = s[i] - margin;
            if(~next(pos, k)) {
                pos = next(pos, k);
                continue;
            }
            const int npos = st.size();
            next(pos, k) = npos;
            st.emplace_back(s[i], pos);
            pos = npos;
        }
        ++st[pos].count;
        st[pos].idx = x;
        st[pos].idxs.emplace_back(x);
    }
    int find(const string& s) {
        int pos = 0;
        for(int i = 0; i < (int)s.size(); ++i) {
            const int k = s[i] - margin;
            if(next(pos, k) < 0) return -1;
            pos = next(pos, k);
        }
        return pos;
    }
    int move(const int pos, const char c) {
        assert(0 <= pos and pos < (int)st.size());
        return next(pos, c - margin);
    }
    int size() const {
        return st.size();
    }
    int idx(const int pos) const {
        assert(0 <= pos and pos < (int)st.size());
        return st[pos].idx;
    }
    int count(const int pos) const {
        assert(0 <= pos and pos < (int)st.size());
        return st[pos].count;
    }
    int par(const int pos) const {
        assert(0 <= pos and pos < (int)st.size());
        return st[pos].parent;
    }
    vector<int> idxs(const int pos) const {
        assert(0 <= pos and pos < (int)st.size());
        return st[pos].idxs;
    }
};

Back to top page