Cod sursa(job #3242427)

Utilizator Mihai145Oprea Mihai Adrian Mihai145 Data 12 septembrie 2024 09:58:54
Problema Aho-Corasick Scor 0
Compilator cpp-64 Status done
Runda Arhiva educationala Marime 3.07 kb
#include <fstream>
#include <string>
#include <vector>
#include <queue>

const int ALPHABET_SIZE = 26;

struct TrieNode
{
    int cnt, fail;
    std::vector<int> to;

    TrieNode()
    {
        cnt = 0, fail = -1;
        to = std::vector<int>(ALPHABET_SIZE, -1);
    }
};

class AhoCorasick
{
private:
    std::vector<TrieNode> nodes;
    std::vector<int> word_ptrs;
    std::vector<int> bfs_order;

public:
    AhoCorasick()
    {
        nodes.push_back({}); // root node
    }

    void insert(const std::string &word)
    {
        int idx = 0;
        for (char c : word)
        {
            if (nodes[idx].to[c - 'a'] != -1)
                idx = nodes[idx].to[c - 'a'];
            else
            {
                nodes.push_back({}); // insert a new node in the trie
                int new_idx = (int)nodes.size() - 1;

                nodes[idx].to[c - 'a'] = new_idx;
                idx = new_idx;
            }
        }

        word_ptrs.push_back(idx);
    }

    void build()
    {
        std::queue<std::tuple<int, int, int>> q;
        q.push({0, -1, 0});

        while (!q.empty())
        {
            auto [idx, parent, chr] = q.front();
            q.pop();
            bfs_order.push_back(idx);

            // calculate fail link
            if (idx == 0)
                nodes[0].fail = 0;
            else
            {
                int c_fail = nodes[parent].fail;
                while (c_fail > 0 && nodes[c_fail].to[chr] == -1)
                    c_fail = nodes[c_fail].fail;

                if (c_fail != parent && nodes[c_fail].to[chr] != -1)
                    c_fail = nodes[c_fail].to[chr];

                nodes[idx].fail = c_fail;
            }

            for (int i = 0; i < ALPHABET_SIZE; i++)
            {
                if (nodes[idx].to[i] == -1)
                    continue;
                q.push({nodes[idx].to[i], idx, i});
            }
        }
    }

    std::vector<int> walk(const std::string &text)
    {
        int idx = 0;
        for (char c : text)
        {
            while (idx > 0 && nodes[idx].to[c - 'a'] == -1)
                idx = nodes[idx].fail;

            if (nodes[idx].to[c - 'a'] != -1)
                idx = nodes[idx].to[c - 'a'];

            nodes[idx].cnt++;
        }

        for (int i = (int)bfs_order.size() - 1; i >= 0; i--)
        {
            int fail = nodes[bfs_order[i]].fail;
            nodes[fail].cnt += nodes[bfs_order[i]].cnt;
        }

        std::vector<int> appearances((int)word_ptrs.size(), 0);
        for (int i = 0; i < (int)word_ptrs.size(); i++)
            appearances[i] = nodes[word_ptrs[i]].cnt;

        return appearances;
    }
};

int main()
{
    std::ifstream cin("ahocorasick.in");
    std::ofstream cout("ahocorasick.out");

    std::string text;
    cin >> text;
    int n;
    cin >> n;

    AhoCorasick automaton;
    for (int i = 0; i < n; i++)
    {
        std::string word;
        cin >> word;
        automaton.insert(word);
    }

    automaton.build();

    std::vector<int> appearances = automaton.walk(text);
    for (int ap : appearances)
        cout << ap << '\n';

    return 0;
}