cpp_src/string/AhoCorasick.hpp

• View this file on GitHub
• Last update: 2023-02-04 17:43:05+09:00
• Include: #include "cpp_src/string/AhoCorasick.hpp"
• Link: https://ei1333.github.io/library/string/aho-corasick.hpp.html

Code

// ref: https://ei1333.github.io/library/string/aho-corasick.hpp.html
// depends on Trie
// ABC268H

template <int char_size, int margin>
struct AhoCorasick : public Trie<char_size + 1, margin> {
    using Trie<char_size + 1, margin>::Trie;

    const int FAIL = char_size;
    vector<int> correct;

    void build(bool heavy = true) {
        correct.resize(this->size());
        for (int i = 0; i < this->size(); i++) {
            correct[i] = (int)this->nodes[i].accept.size();
        }
        queue<int> que;
        for (int i = 0; i <= char_size; i++) {
            if (~this->nodes[0].nxt[i]) {
                this->nodes[this->nodes[0].nxt[i]].nxt[FAIL] = 0;
                que.emplace(this->nodes[0].nxt[i]);
            } else {
                this->nodes[0].nxt[i] = 0;
            }
        }
        while (!que.empty()) {
            auto& now = this->nodes[que.front()];
            int fail = now.nxt[FAIL];
            correct[que.front()] += correct[fail];
            que.pop();
            for (int i = 0; i < char_size; i++) {
                if (~now.nxt[i]) {
                    this->nodes[now.nxt[i]].nxt[FAIL] =
                        this->nodes[fail].nxt[i];
                    if (heavy) {
                        auto& u = this->nodes[now.nxt[i]].accept;
                        auto& v = this->nodes[this->nodes[fail].nxt[i]].accept;
                        vector<int> accept;
                        set_union(begin(u), end(u), begin(v), end(v),
                                  back_inserter(accept));
                        u = accept;
                    }
                    que.emplace(now.nxt[i]);
                } else {
                    now.nxt[i] = this->nodes[fail].nxt[i];
                }
            }
        }
    }

    map<int, int> match(const string& str, int now = 0) {
        map<int, int> result;
        for (auto& c : str) {
            now = this->nodes[now].nxt[c - margin];
            for (auto& v : this->nodes[now].accept) result[v] += 1;
        }
        return result;
    }

    pair<int64_t, int> move(const char& c, int now = 0) {
        now = this->nodes[now].nxt[c - margin];
        return {correct[now], now};
    }

    pair<int64_t, int> move(const string& str, int now = 0) {
        int64_t sum = 0;
        for (auto& c : str) {
            auto nxt = move(c, now);
            sum += nxt.first;
            now = nxt.second;
        }
        return {sum, now};
    }
};

#line 1 "cpp_src/string/AhoCorasick.hpp"
// ref: https://ei1333.github.io/library/string/aho-corasick.hpp.html
// depends on Trie
// ABC268H

template <int char_size, int margin>
struct AhoCorasick : public Trie<char_size + 1, margin> {
    using Trie<char_size + 1, margin>::Trie;

    const int FAIL = char_size;
    vector<int> correct;

    void build(bool heavy = true) {
        correct.resize(this->size());
        for (int i = 0; i < this->size(); i++) {
            correct[i] = (int)this->nodes[i].accept.size();
        }
        queue<int> que;
        for (int i = 0; i <= char_size; i++) {
            if (~this->nodes[0].nxt[i]) {
                this->nodes[this->nodes[0].nxt[i]].nxt[FAIL] = 0;
                que.emplace(this->nodes[0].nxt[i]);
            } else {
                this->nodes[0].nxt[i] = 0;
            }
        }
        while (!que.empty()) {
            auto& now = this->nodes[que.front()];
            int fail = now.nxt[FAIL];
            correct[que.front()] += correct[fail];
            que.pop();
            for (int i = 0; i < char_size; i++) {
                if (~now.nxt[i]) {
                    this->nodes[now.nxt[i]].nxt[FAIL] =
                        this->nodes[fail].nxt[i];
                    if (heavy) {
                        auto& u = this->nodes[now.nxt[i]].accept;
                        auto& v = this->nodes[this->nodes[fail].nxt[i]].accept;
                        vector<int> accept;
                        set_union(begin(u), end(u), begin(v), end(v),
                                  back_inserter(accept));
                        u = accept;
                    }
                    que.emplace(now.nxt[i]);
                } else {
                    now.nxt[i] = this->nodes[fail].nxt[i];
                }
            }
        }
    }

    map<int, int> match(const string& str, int now = 0) {
        map<int, int> result;
        for (auto& c : str) {
            now = this->nodes[now].nxt[c - margin];
            for (auto& v : this->nodes[now].accept) result[v] += 1;
        }
        return result;
    }

    pair<int64_t, int> move(const char& c, int now = 0) {
        now = this->nodes[now].nxt[c - margin];
        return {correct[now], now};
    }

    pair<int64_t, int> move(const string& str, int now = 0) {
        int64_t sum = 0;
        for (auto& c : str) {
            auto nxt = move(c, now);
            sum += nxt.first;
            now = nxt.second;
        }
        return {sum, now};
    }
};

Back to top page