Cod sursa(job #3233307)

Utilizator MirceaDonciuLicentaLicenta Mircea Donciu MirceaDonciuLicenta Data 2 iunie 2024 22:51:37
Problema Aho-Corasick Scor 95
Compilator cpp-64 Status done
Runda Arhiva educationala Marime 2.97 kb
#include <iostream>
#include <fstream>
#include <vector>
#include <queue>
#include <unordered_map>

using namespace std;

const int ALPHABET_SIZE = 26;

struct TrieNode {
    vector<TrieNode*> children;
    TrieNode* fail;
    vector<int> output;

    TrieNode() : children(ALPHABET_SIZE, nullptr), fail(nullptr) {}
};

class AhoCorasick {
public:
    AhoCorasick() {
        root = new TrieNode();
    }

    void insert(const string& word, int index) {
        TrieNode* node = root;
        for (char ch : word) {
            int c = ch - 'a';
            if (!node->children[c]) {
                node->children[c] = new TrieNode();
            }
            node = node->children[c];
        }
        node->output.push_back(index);
    }

    void build() {
        queue<TrieNode*> q;
        root->fail = root;
        for (int c = 0; c < ALPHABET_SIZE; ++c) {
            if (root->children[c]) {
                root->children[c]->fail = root;
                q.push(root->children[c]);
            } else {
                root->children[c] = root;
            }
        }

        while (!q.empty()) {
            TrieNode* current = q.front();
            q.pop();

            for (int c = 0; c < ALPHABET_SIZE; ++c) {
                if (current->children[c]) {
                    TrieNode* fail = current->fail;
                    while (!fail->children[c]) {
                        fail = fail->fail;
                    }
                    current->children[c]->fail = fail->children[c];
                    current->children[c]->output.insert(
                        current->children[c]->output.end(),
                        current->children[c]->fail->output.begin(),
                        current->children[c]->fail->output.end()
                    );
                    q.push(current->children[c]);
                }
            }
        }
    }

    vector<int> search(const string& text, int numWords) {
        vector<int> counts(numWords, 0);
        TrieNode* node = root;

        for (char ch : text) {
            int c = ch - 'a';
            while (!node->children[c]) {
                node = node->fail;
            }
            node = node->children[c];

            for (int index : node->output) {
                counts[index]++;
            }
        }

        return counts;
    }

private:
    TrieNode* root;
};

int main() {
    ifstream infile("ahocorasick.in");
    ofstream outfile("ahocorasick.out");

    string text;
    infile >> text;

    int numWords;
    infile >> numWords;

    AhoCorasick aho;
    vector<string> words(numWords);

    for (int i = 0; i < numWords; ++i) {
        infile >> words[i];
        aho.insert(words[i], i);
    }

    aho.build();

    vector<int> results = aho.search(text, numWords);

    for (int count : results) {
        outfile << count << endl;
    }

    infile.close();
    outfile.close();

    return 0;
}