string/aho_corasick.hpp

• View this file on GitHub
• Last update: 2022-07-30 11:55:56+09:00
• Include: #include "string/aho_corasick.hpp"

Depends on

• string/trie.hpp

Code

#pragma once
#include <queue>

#include "trie.hpp"

std::vector<std::size_t> aho_corasick(const Trie &trie) {
    std::queue<std::size_t> que;
    std::vector<std::size_t> fail(trie.size(), 0);
    for (std::size_t i = 0; i < 26; ++i) {
        if (trie[0][i] != Trie::none) {
            fail[trie[0][i]] = 0;
            que.push(trie[0][i]);
        }
    }
    while (!que.empty()) {
        std::size_t v = que.front();
        que.pop();
        for (std::size_t i = 0; i < 26; ++i) {
            if (trie[v][i] == Trie::none) {
                continue;
            }
            std::size_t cur = fail[v];
            while (cur != 0 && trie[cur][i] == Trie::none) {
                cur = fail[cur];
            }
            if (trie[cur][i] == Trie::none) {
                fail[trie[v][i]] = 0;
            } else {
                fail[trie[v][i]] = trie[cur][i];
            }
            que.push(trie[v][i]);
        }
    }
    return fail;
}

#line 2 "string/aho_corasick.hpp"
#include <queue>

#line 2 "string/trie.hpp"
#include <algorithm>
#include <array>
#include <string>
#include <vector>

class Trie {
    std::vector<std::array<int, 26>> trie;
    std::vector<int> cnt;

public:
    static constexpr int none = -1;

private:
    int make_node() {
        int s = trie.size();
        trie.emplace_back(std::array<int, 26>());
        std::fill(trie[s].begin(), trie[s].end(), none);
        cnt.push_back(0);
        return s;
    }

public:
    Trie() : trie(), cnt() {
        make_node();
    }

    int size() const {
        return (int) trie.size();
    }

    const std::array<int, 26> &operator[](int i) const {
        return trie[i];
    }

    int insert(const std::string &s) {
        int cur = 0;
        for (char c : s) {
            if (trie[cur][c - 'a'] == none) {
                trie[cur][c - 'a'] = make_node();
            }
            cur = trie[cur][c - 'a'];
        }
        return cnt[cur]++;
    }
    
    int find(const std::string &s) const {
        int cur = 0;
        for (char c : s) {
            int nxt = trie[cur][c - 'a'];
            if (nxt == none) {
                return none;
            }
            cur = nxt;
        }
        return cur;
    }

    int count(const std::string &s) const {
        int node = find(s);
        if (node == none) {
            return 0;
        } else {
            return cnt[node];
        }
    }

    int cnt_node(int i) const {
        return cnt[i];
    }
};
#line 5 "string/aho_corasick.hpp"

std::vector<std::size_t> aho_corasick(const Trie &trie) {
    std::queue<std::size_t> que;
    std::vector<std::size_t> fail(trie.size(), 0);
    for (std::size_t i = 0; i < 26; ++i) {
        if (trie[0][i] != Trie::none) {
            fail[trie[0][i]] = 0;
            que.push(trie[0][i]);
        }
    }
    while (!que.empty()) {
        std::size_t v = que.front();
        que.pop();
        for (std::size_t i = 0; i < 26; ++i) {
            if (trie[v][i] == Trie::none) {
                continue;
            }
            std::size_t cur = fail[v];
            while (cur != 0 && trie[cur][i] == Trie::none) {
                cur = fail[cur];
            }
            if (trie[cur][i] == Trie::none) {
                fail[trie[v][i]] = 0;
            } else {
                fail[trie[v][i]] = trie[cur][i];
            }
            que.push(trie[v][i]);
        }
    }
    return fail;
}

Back to top page