test/unit/aho_corasick.test.cpp

View this file on GitHub
Last update: 2026-10-27 07:43:44+09:00
Problem: http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=ITP1_1_A

Depends on

Code

#define PROBLEM "http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=ITP1_1_A"
#include <algorithm>
#include <cassert>
#include <iostream>
#include <random>
#include <string>
#include <vector>

#include "string/aho_corasick.hpp"
#include "string/trie.hpp"
std::string rand_str(std::mt19937& rng, int maxlen, int alpha) {
    int len = rng() % (maxlen + 1);
    std::string s;
    for (int i = 0; i < len; i++) s += char('a' + rng() % alpha);
    return s;
}
bool ends_with(const std::string& s, const std::string& p) {
    return p.size() <= s.size() && s.compare(s.size() - p.size(), p.size(), p) == 0;
}
int main() {
    std::mt19937 rng(430);
    for (int it = 0; it < 500; it++) {
        int alpha = 1 + rng() % 3, m = rng() % 8;
        trie<3> t;
        std::vector<std::string> pats;
        std::vector<int> end;
        for (int i = 0; i < m; i++) {
            pats.push_back(rand_str(rng, 4, alpha));
            end.push_back(t.insert(pats.back()));
        }
        aho_corasick ac(t);
        assert(ac.size() == t.size());
        // order: increasing depth, starts with root
        const auto& ord = ac.order();
        assert(int(ord.size()) == t.size() && ord[0] == 0);
        for (int i = 1; i < int(ord.size()); i++) assert(t.depth(ord[i - 1]) <= t.depth(ord[i]));
        // fail: longest proper suffix that is a trie node
        std::vector<std::string> str(t.size());
        for (int v : ord) {
            if (v == 0) continue;
            int p = t.parent(v);
            for (int c = 0; c < 3; c++) {
                if (t.next(p, 'a' + c) == v) str[v] = str[p] + char('a' + c);
            }
        }
        for (int v = 0; v < t.size(); v++) {
            if (v == 0) {
                assert(ac.fail(v) == -1);
                continue;
            }
            int best = 0;
            for (int u = 0; u < t.size(); u++) {
                if (u != v && ends_with(str[v], str[u]) && str[u].size() > str[best].size()) best = u;
            }
            assert(ac.fail(v) == best);
            int c = 0;
            for (auto& p : pats) c += ends_with(str[v], p);
            assert(ac.count(v) == c);
        }
        // scan a text: state = longest suffix of the read prefix that is a trie node
        std::string s = rand_str(rng, 30, alpha);
        std::vector<long long> hit(t.size());
        std::vector<int> occ(m);
        int v = 0;
        long long total = 0;
        for (int i = 0; i < int(s.size()); i++) {
            v = ac.next(v, s[i]);
            hit[v]++;
            total += ac.count(v);
            std::string pre = s.substr(0, i + 1);
            int best = 0;
            for (int u = 0; u < t.size(); u++) {
                if (ends_with(pre, str[u]) && str[u].size() > str[best].size()) best = u;
            }
            assert(v == best);
            for (int j = 0; j < m; j++) occ[j] += ends_with(pre, pats[j]);
        }
        long long brute = 0;
        for (int j = 0; j < m; j++) brute += occ[j];
        assert(total == brute);
        for (int i = int(ord.size()) - 1; i >= 1; i--) hit[ac.fail(ord[i])] += hit[ord[i]];
        for (int j = 0; j < m; j++) assert(hit[end[j]] == occ[j]);
    }
    {
        // empty pattern matches at every position
        trie<2> t;
        t.insert(std::string());
        aho_corasick ac(t);
        assert(ac.size() == 1 && ac.count(0) == 1 && ac.next(0, 'a') == 0 && ac.next(0, 'b') == 0);
    }
    {
        // large: 2e5 patterns of length 5 over 26 letters, text 2e5
        trie<26> t;
        for (int i = 0; i < 200000; i++) t.insert(rand_str(rng, 5, 26) + "a");
        aho_corasick ac(t);
        int v = 0;
        long long total = 0;
        for (int i = 0; i < 200000; i++) {
            v = ac.next(v, 'a' + rng() % 26);
            total += ac.count(v);
        }
        assert(total >= 0);
    }
    std::cout << "Hello World" << std::endl;
}

#line 1 "test/unit/aho_corasick.test.cpp"
#define PROBLEM "http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=ITP1_1_A"
#include <algorithm>
#include <cassert>
#include <iostream>
#include <random>
#include <string>
#include <vector>

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

#line 5 "string/trie.hpp"
template <int sigma = 26, int base = 'a'> struct trie {
    static_assert(sigma >= 1);

public:
    trie() : ch(1), par(1, -1), dep(1, 0), cnt(1, 0), cnt_prefix(1, 0) {
        ch[0].fill(-1);
    }
    int size() const {
        return int(ch.size());
    }
    template <class Seq> int insert(const Seq& s) {
        int v = 0;
        cnt_prefix[v]++;
        for (auto c : s) {
            int x = int(c) - base;
            assert(0 <= x && x < sigma);
            if (ch[v][x] == -1) {
                ch[v][x] = size();
                ch.emplace_back().fill(-1);
                par.push_back(v);
                dep.push_back(dep[v] + 1);
                cnt.push_back(0);
                cnt_prefix.push_back(0);
            }
            v = ch[v][x];
            cnt_prefix[v]++;
        }
        cnt[v]++;
        return v;
    }
    template <class Seq> void erase(const Seq& s) {
        int t = find(s);
        assert(t != -1 && cnt[t] > 0);
        cnt[t]--;
        for (int v = t; v != -1; v = par[v]) cnt_prefix[v]--;
    }
    template <class Seq> int find(const Seq& s) const {
        int v = 0;
        for (auto c : s) {
            v = next(v, int(c));
            if (v == -1) return -1;
        }
        return v;
    }
    int next(int v, int c) const {
        assert(0 <= v && v < size());
        int x = c - base;
        assert(0 <= x && x < sigma);
        return ch[v][x];
    }
    int parent(int v) const {
        assert(0 <= v && v < size());
        return par[v];
    }
    int depth(int v) const {
        assert(0 <= v && v < size());
        return dep[v];
    }
    int count(int v) const {
        assert(0 <= v && v < size());
        return cnt[v];
    }
    int count_prefix(int v) const {
        assert(0 <= v && v < size());
        return cnt_prefix[v];
    }

private:
    std::vector<std::array<int, sigma>> ch;
    std::vector<int> par, dep, cnt, cnt_prefix;
};
#line 7 "string/aho_corasick.hpp"
template <int sigma = 26, int base = 'a'> struct aho_corasick {
    explicit aho_corasick(const trie<sigma, base>& t) : _n(t.size()), nxt(_n), fl(_n, -1), cnt(_n) {
        ord.reserve(_n);
        ord.push_back(0);
        for (int i = 0; i < _n; i++) {
            int v = ord[i];
            cnt[v] = t.count(v) + (v == 0 ? 0 : cnt[fl[v]]);
            for (int c = 0; c < sigma; c++) {
                int w = t.next(v, base + c);
                if (w == -1) {
                    nxt[v][c] = v == 0 ? 0 : nxt[fl[v]][c];
                } else {
                    nxt[v][c] = w;
                    fl[w] = v == 0 ? 0 : nxt[fl[v]][c];
                    ord.push_back(w);
                }
            }
        }
        assert(int(ord.size()) == _n);
    }
    int size() const {
        return _n;
    }
    int next(int v, int c) const {
        assert(0 <= v && v < _n);
        int x = c - base;
        assert(0 <= x && x < sigma);
        return nxt[v][x];
    }
    int fail(int v) const {
        assert(0 <= v && v < _n);
        return fl[v];
    }
    int count(int v) const {
        assert(0 <= v && v < _n);
        return cnt[v];
    }
    const std::vector<int>& order() const {
        return ord;
    }

  private:
    int _n;
    std::vector<std::array<int, sigma>> nxt;
    std::vector<int> fl, cnt, ord;
};
#line 11 "test/unit/aho_corasick.test.cpp"
std::string rand_str(std::mt19937& rng, int maxlen, int alpha) {
    int len = rng() % (maxlen + 1);
    std::string s;
    for (int i = 0; i < len; i++) s += char('a' + rng() % alpha);
    return s;
}
bool ends_with(const std::string& s, const std::string& p) {
    return p.size() <= s.size() && s.compare(s.size() - p.size(), p.size(), p) == 0;
}
int main() {
    std::mt19937 rng(430);
    for (int it = 0; it < 500; it++) {
        int alpha = 1 + rng() % 3, m = rng() % 8;
        trie<3> t;
        std::vector<std::string> pats;
        std::vector<int> end;
        for (int i = 0; i < m; i++) {
            pats.push_back(rand_str(rng, 4, alpha));
            end.push_back(t.insert(pats.back()));
        }
        aho_corasick ac(t);
        assert(ac.size() == t.size());
        // order: increasing depth, starts with root
        const auto& ord = ac.order();
        assert(int(ord.size()) == t.size() && ord[0] == 0);
        for (int i = 1; i < int(ord.size()); i++) assert(t.depth(ord[i - 1]) <= t.depth(ord[i]));
        // fail: longest proper suffix that is a trie node
        std::vector<std::string> str(t.size());
        for (int v : ord) {
            if (v == 0) continue;
            int p = t.parent(v);
            for (int c = 0; c < 3; c++) {
                if (t.next(p, 'a' + c) == v) str[v] = str[p] + char('a' + c);
            }
        }
        for (int v = 0; v < t.size(); v++) {
            if (v == 0) {
                assert(ac.fail(v) == -1);
                continue;
            }
            int best = 0;
            for (int u = 0; u < t.size(); u++) {
                if (u != v && ends_with(str[v], str[u]) && str[u].size() > str[best].size()) best = u;
            }
            assert(ac.fail(v) == best);
            int c = 0;
            for (auto& p : pats) c += ends_with(str[v], p);
            assert(ac.count(v) == c);
        }
        // scan a text: state = longest suffix of the read prefix that is a trie node
        std::string s = rand_str(rng, 30, alpha);
        std::vector<long long> hit(t.size());
        std::vector<int> occ(m);
        int v = 0;
        long long total = 0;
        for (int i = 0; i < int(s.size()); i++) {
            v = ac.next(v, s[i]);
            hit[v]++;
            total += ac.count(v);
            std::string pre = s.substr(0, i + 1);
            int best = 0;
            for (int u = 0; u < t.size(); u++) {
                if (ends_with(pre, str[u]) && str[u].size() > str[best].size()) best = u;
            }
            assert(v == best);
            for (int j = 0; j < m; j++) occ[j] += ends_with(pre, pats[j]);
        }
        long long brute = 0;
        for (int j = 0; j < m; j++) brute += occ[j];
        assert(total == brute);
        for (int i = int(ord.size()) - 1; i >= 1; i--) hit[ac.fail(ord[i])] += hit[ord[i]];
        for (int j = 0; j < m; j++) assert(hit[end[j]] == occ[j]);
    }
    {
        // empty pattern matches at every position
        trie<2> t;
        t.insert(std::string());
        aho_corasick ac(t);
        assert(ac.size() == 1 && ac.count(0) == 1 && ac.next(0, 'a') == 0 && ac.next(0, 'b') == 0);
    }
    {
        // large: 2e5 patterns of length 5 over 26 letters, text 2e5
        trie<26> t;
        for (int i = 0; i < 200000; i++) t.insert(rand_str(rng, 5, 26) + "a");
        aho_corasick ac(t);
        int v = 0;
        long long total = 0;
        for (int i = 0; i < 200000; i++) {
            v = ac.next(v, 'a' + rng() % 26);
            total += ac.count(v);
        }
        assert(total >= 0);
    }
    std::cout << "Hello World" << std::endl;
}