Cod sursa(job #3124796)

Utilizator Vlad_NituNitu Vlad-Petru Vlad_Nitu Data 30 aprilie 2023 02:06:03
Problema Aho-Corasick Scor 100
Compilator cpp-64 Status done
Runda Arhiva educationala Marime 2.92 kb
#include <bits/stdc++.h>

using namespace std;

const int ALPHABET_SIZE = 26;

struct Aho {
    Aho *children[ALPHABET_SIZE];
    Aho *failure;
    vector<int> words;
    int cnt;

    Aho() {
        cnt = 0;
        failure = nullptr;
        words = {};
        for (int i = 0; i < ALPHABET_SIZE; ++i)
            children[i] = nullptr;
    }
};

Aho *AHC = new Aho();

// construct Trie
// word_id - index of word in vector<string>& w
// pos - ch. position in w[word_id]
void insert(Aho *node, int pos, const int &word_id, const vector<string> &w) {
    if (pos == w[word_id].size()) {
        node->words.push_back(word_id);
        return;
    }

    const string &word = w[word_id];
    if (node->children[word[pos] - 'a'] == nullptr)
        node->children[word[pos] - 'a'] = new Aho();
    insert(node->children[word[pos] - 'a'], pos + 1, word_id, w);
}

void BFS() {

    AHC->failure = AHC;
    queue<Aho *> Q;

    // All nodes 's' of depth 1 have "failure link" back to root
    for (int i = 0; i < ALPHABET_SIZE; ++i)
        if (AHC->children[i]) {
            AHC->children[i]->failure = AHC;
            Q.push(AHC->children[i]);
        }

    while (!Q.empty()) {
        Aho *curr = Q.front();
        Q.pop();

        for (int i = 0; i < ALPHABET_SIZE; ++i)
            if (curr->children[i]) {
                Aho *fail = curr->failure;

                while (fail != AHC && fail->children[i] == nullptr)
                    fail = fail->failure;

                if (fail->children[i] && fail->children[i] != curr->children[i])
                    fail = fail->children[i];

                curr->children[i]->failure = fail;
                Q.push(curr->children[i]);
            }
    }
}

vector<int> scanText(const int &K, const string &text) {

    vector<int> freq(K);
    Aho *head = AHC;

    for (const auto& ch: text) {
        while (head != AHC && head->children[ch - 'a'] == nullptr)
            head = head->failure;
        if (head->children[ch - 'a'])
            head = head->children[ch - 'a'];

        ++head->cnt;
    }

    queue<Aho *> Q;
    vector<Aho *> nodes;
    Q.push(AHC);

    while (!Q.empty()) {
        Aho *curr = Q.front();
        Q.pop();

        nodes.push_back(curr);
        for (int i = 0; i < ALPHABET_SIZE; ++i)
            if (curr->children[i])
                Q.push(curr->children[i]);
    }

    reverse(nodes.begin(), nodes.end());

    for (const Aho *node: nodes) {

        for (auto word_id: node->words)
            freq[word_id] += node->cnt;

        node->failure->cnt += node->cnt;
    }
    return freq;
}

string text;
int K;

int main() {

    freopen("ahocorasick.in", "r", stdin);
    freopen("ahocorasick.out", "w", stdout);

    cin >> text;
    cin >> K;
    vector<string> w(K);
    for (int i = 0; i < K; ++i) {
        cin >> w[i];
        insert(AHC, 0, i, w);
    }


    BFS();

    vector<int> ans = scanText(K, text);
    for (const auto& fq: ans)
        cout << fq << "\n";

    return 0;
}