cp-documentation

This documentation is automatically generated by online-judge-tools/verification-helper

View the Project on GitHub zawa-tin/cp-documentation

:heavy_check_mark: Test/LC/aho_corasick.test.cpp

Depends on

Code

#define PROBLEM "https://judge.yosupo.jp/problem/aho_corasick"

#include "../../Src/Sequence/AhoCorasick.hpp"

#include <iostream>
#include <string>
#include <vector>

using namespace std;
int main() {
    cin.tie(0);
    cout.tie(0);
    ios::sync_with_stdio(0);
    int N;
    cin >> N;
    zawa::AhoCorasick<string> aho;
    for (int i = 0 ; i < N ; i++) {
        string s;
        cin >> s;
        aho.insert(s);
    }
    auto trie = aho.build();
    const auto& nodes = trie.nodes(); 
    cout << trie.size() << '\n';
    for (int i = 1 ; i < ssize(nodes) ; i++)
        cout << nodes[i].par.first << ' ' << nodes[i].fail << '\n';
    for (int i = 0 ; i < N ; i++)
        cout << trie.match(i) << (i + 1 == N ? '\n' : ' ');
}
#line 1 "Test/LC/aho_corasick.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/aho_corasick"

#line 2 "Src/Sequence/AhoCorasick.hpp"

#line 2 "Src/Template/TypeAlias.hpp"

#include <cstdint>
#include <cstddef>

namespace zawa {

using i16 = std::int16_t;
using i32 = std::int32_t;
using i64 = std::int64_t;
using i128 = __int128_t;

using u8 = std::uint8_t;
using u16 = std::uint16_t;
using u32 = std::uint32_t;
using u64 = std::uint64_t;

using usize = std::size_t;

} // namespace zawa
#line 2 "Src/Algebra/Monoid/MonoidConcept.hpp"

#line 2 "Src/Algebra/Semigroup/SemigroupConcept.hpp"

#include <concepts>

namespace zawa {

namespace concepts {

template <class T>
concept Semigroup = requires {
    typename T::Element;
    { T::operation(std::declval<typename T::Element>(), std::declval<typename T::Element>()) } -> std::same_as<typename T::Element>;
};

} // namespace concepts

} // namespace zawa
#line 4 "Src/Algebra/Monoid/MonoidConcept.hpp"

#line 6 "Src/Algebra/Monoid/MonoidConcept.hpp"

namespace zawa {

namespace concepts {

template <class T>
concept Identitiable = requires {
    typename T::Element;
    { T::identity() } -> std::same_as<typename T::Element>;
};

template <class T>
concept Monoid = Semigroup<T> and Identitiable<T>;

} // namespace

} // namespace zawa
#line 2 "Src/Algebra/Action/ActionConcept.hpp"

#line 4 "Src/Algebra/Action/ActionConcept.hpp"

namespace zawa {

namespace concepts {

template <class G, class X>
concept Action = requires {
    typename G::Element;
    { G::action(std::declval<typename G::Element>(), std::declval<X>()) } -> std::same_as<X>;
};

// Is appropriate name X-set?
template <class G, class X>
concept Acted = requires {
    typename G::Element;
    { G::acted(std::declval<typename G::Element>(), std::declval<X>()) } -> std::same_as<typename G::Element>;
};

} // namespace concepts

} // namespace zawa
#line 6 "Src/Sequence/AhoCorasick.hpp"

#include <cassert>
#line 9 "Src/Sequence/AhoCorasick.hpp"
#include <ranges>
#include <unordered_map>
#include <utility>
#include <vector>

namespace zawa {

namespace ahocorasick_internal {

template <class T>
concept HasValueType = requires {
    typename T::value_type;
};

template <class T, class S>
concept AuxiliaryData = concepts::Monoid<T> and concepts::Acted<T, S>;

} // namespace ahocorasick_internal

template <ahocorasick_internal::HasValueType Container>
class AhoCorasick {
public:

    using V = Container::value_type;

private:

    class Trie {
    public:

        struct Node {
            usize fail = 0;
            std::unordered_map<V, usize> ch{};
            std::pair<usize, V> par{};
        };

        Trie(std::vector<Node>&& nodes, std::vector<usize>&& match) 
            : m_nodes{std::move(nodes)}, m_match{std::move(match)} {}

        static constexpr usize Root() {
            return 0;
        }

        usize size() const {
            return m_nodes.size();
        }

        usize trace(usize cur, V v) {
            assert(cur < size());
            while (cur and !m_nodes[cur].ch.contains(v))
                cur = m_nodes[cur].fail;
            if (auto it = m_nodes[cur].ch.find(v) ; it != m_nodes[cur].ch.end())
                return it->second;
            else
                return cur;
        }

        usize match(usize i) const {
            assert(i < m_match.size());
            return m_match[i];
        }

        usize trace(usize cur, const Container& S) {
            assert(cur < size());
            for (V v : S)
                cur = trace(cur, v);
            return cur;
        }

        const std::vector<Node>& nodes() const& {
            return m_nodes;
        }

    private:

        std::vector<Node> m_nodes;

        std::vector<usize> m_match;
    };

public:

    AhoCorasick() = default;

    usize insert(Container s) {
        usize res = m_seq.size();
        m_seq.push_back(s);
        return res;
    }

    Trie build() const {
        std::vector<typename Trie::Node> nodes(1);  
        std::vector<usize> match;
        match.reserve(m_seq.size());
        for (const Container& s : m_seq) {
            usize cur = 0, idx = 0;
            for ( ; idx < s.size() ; idx++) {
                auto it = nodes[cur].ch.find(s[idx]);
                if (it == nodes[cur].ch.end())
                    break;
                cur = it->second;
            }
            for ( ; idx < s.size() ; idx++) {
                usize nxt = nodes[cur].ch[s[idx]] = nodes.size();
                nodes.emplace_back();
                nodes.back().par = {cur, s[idx]};
                cur = nxt;
            }
            match.push_back(cur);
        }
        std::vector<usize> que;
        for (const usize x : nodes[0].ch | std::views::values)
            que.emplace_back(x);
        for (usize qt = 0 ; qt < que.size() ; qt++) {
            const usize v = que[qt];
            for (const usize x : nodes[v].ch | std::views::values)
                que.emplace_back(x);
            auto [x, ed] = nodes[v].par;
            if (!x) 
                continue;
            x = nodes[x].fail;
            while (x and !nodes[x].ch.contains(ed))
                x = nodes[x].fail;
            if (auto it = nodes[x].ch.find(ed) ; it == nodes[x].ch.end() or it->second == v) 
                nodes[v].fail = 0;
            else
                nodes[v].fail = it->second;
        }
        return Trie{std::move(nodes), std::move(match)};
    }

    template <class T, class S>
    requires ahocorasick_internal::AuxiliaryData<T, S>
    std::pair<Trie, std::vector<typename T::Element>> build(const std::vector<S>& values) const {
        assert(values.size() == m_seq.size());
        std::vector<typename T::Element> data(1, T::identity());
        std::vector<typename Trie::Node> nodes(1);  
        std::vector<usize> match(m_seq.size());
        for (usize i = 0 ; const Container& s : m_seq) {
            usize cur = 0, idx = 0;
            for ( ; idx < s.size() ; idx++) {
                auto it = nodes[cur].ch.find(s[idx]);
                if (it == nodes[cur].ch.end())
                    break;
                cur = it->second;
            }
            for ( ; idx < s.size() ; idx++) {
                usize nxt = nodes[cur].ch[s[idx]] = nodes.size();
                nodes.emplace_back();
                nodes.back().par = {cur, s[idx]};
                data.push_back(data[cur]);
                cur = nxt;
            }
            match[i] = cur;
            data[cur] = T::acted(data[cur], values[i++]);
        }
        std::vector<usize> que;
        for (const usize x : nodes[0].ch | std::views::values)
            que.emplace_back(x);
        for (usize qt = 0 ; qt < que.size() ; qt++) {
            const usize v = que[qt];
            for (const usize x : nodes[v].ch | std::views::values)
                que.emplace_back(x);
            auto [x, ed] = nodes[v].par;
            if (!x) 
                continue;
            x = nodes[x].fail;
            while (x and !nodes[x].ch.contains(ed))
                x = nodes[x].fail;
            if (auto it = nodes[x].ch.find(ed) ; it == nodes[x].ch.end() or it->second == v) 
                nodes[v].fail = 0;
            else
                nodes[v].fail = it->second;
            data[v] = T::operation(data[nodes[v].fail], data[v]);
        }
        return std::pair{Trie{std::move(nodes), std::move(match)}, data};
    }

private:

    std::vector<Container> m_seq;

};

} // namespace zawa
#line 4 "Test/LC/aho_corasick.test.cpp"

#include <iostream>
#include <string>
#line 8 "Test/LC/aho_corasick.test.cpp"

using namespace std;
int main() {
    cin.tie(0);
    cout.tie(0);
    ios::sync_with_stdio(0);
    int N;
    cin >> N;
    zawa::AhoCorasick<string> aho;
    for (int i = 0 ; i < N ; i++) {
        string s;
        cin >> s;
        aho.insert(s);
    }
    auto trie = aho.build();
    const auto& nodes = trie.nodes(); 
    cout << trie.size() << '\n';
    for (int i = 1 ; i < ssize(nodes) ; i++)
        cout << nodes[i].par.first << ' ' << nodes[i].fail << '\n';
    for (int i = 0 ; i < N ; i++)
        cout << trie.match(i) << (i + 1 == N ? '\n' : ' ');
}
Back to top page