< ゾビオさんのAI生成コード置き場

Trie木

日時はサーバーが記録したものです。更新すると更新日時が変わるので、不都合がある場合は新しく投稿してください。

struct Trie {
    long long K;              // 文字種数
    char base;                // 先頭文字（'a' など）
    vector<vector<long long>> next;
    vector<long long> cnt;    // そのノードを通過した単語数
    vector<long long> end;    // そのノードで終わる単語数

    Trie(long long k = 26, char b = 'a') {
        K = k;
        base = b;
        add_node();           // 根を作る
    }

    long long add_node() {
        next.push_back(vector<long long>(K, -1));
        cnt.push_back(0);
        end.push_back(0);
        return (long long)next.size() - 1;
    }

    void insert(const string &s) {
        long long v = 0;
        for (char ch : s) {
            long long c = ch - base;
            if (next.at(v).at(c) == -1)
                next[v][c] = add_node();
            v = next[v][c];
            cnt[v]++;
        }
        end[v]++;
    }

    void erase(const string &s) {  // 存在する前提
        long long v = 0;
        for (char ch : s) {
            v = next[v][ch - base];
            cnt[v]--;
        }
        end[v]--;
    }

    // s を辿った先のノード番号（無ければ -1）
    long long find(const string &s) {
        long long v = 0;
        for (char ch : s) {
            v = next.at(v).at(ch - base);
            if (v == -1)
                return -1;
        }
        return v;
    }

    // s と完全一致する単語の個数
    long long count(const string &s) {
        long long v = find(s);
        return v == -1 ? 0 : end.at(v);
    }

    // s を接頭辞に持つ単語の個数
    long long count_prefix(const string &s) {
        long long v = find(s);
        return v == -1 ? 0 : cnt.at(v);
    }

    bool contains(const string &s) {
        return count(s) > 0;
    }

    // s の接頭辞のうち、辞書に入っている単語の長さ一覧
    vector<long long> common_prefix_lengths(const string &s) {
        vector<long long> ret;
        long long v = 0;
        for (long long i = 0; i < (long long)s.size(); i++) {
            v = next.at(v).at(s[i] - base);
            if (v == -1)
                break;
            if (end[v] > 0)
                ret.push_back(i + 1);
        }
        return ret;
    }

    // s 以下（辞書順）の文字列の個数  ※全て同じ長さである前提
    long long count_less_equal(const string &s) {
        long long v = 0, res = 0;
        for (char ch : s) {
            long long c = ch - base;
            // s の現在の文字より小さい文字へ分岐したものは全て s より小さい
            for (long long d = 0; d < c; d++)
                if (next.at(v).at(d) != -1)
                    res += cnt[next[v][d]];
            if (next[v][c] == -1)
                return res;      // これ以上 s と一致する枝がない
            v = next[v][c];
        }
        res += cnt[v];           // s と一致するもの（同じ長さ前提なので cnt でよい）
        return res;
    }

    long long node_count() {
        return (long long)next.size();
    }
};