Cod sursa(job #3233306)

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

using namespace std;

struct TrieNode {
    unordered_map<char, TrieNode*> children;
    TrieNode* fail;
    vector<int> output;
};

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

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

    void build() {
        queue<TrieNode*> q;
        root->fail = root;
        for (auto& pair : root->children) {
            pair.second->fail = root;
            q.push(pair.second);
        }

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

            for (auto& pair : current->children) {
                char ch = pair.first;
                TrieNode* child = pair.second;

                TrieNode* fail = current->fail;
                while (fail != root && fail->children.find(ch) == fail->children.end()) {
                    fail = fail->fail;
                }

                if (fail->children.find(ch) != fail->children.end()) {
                    child->fail = fail->children[ch];
                } else {
                    child->fail = root;
                }

                child->output.insert(child->output.end(), child->fail->output.begin(), child->fail->output.end());
                q.push(child);
            }
        }
    }

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

        for (char ch : text) {
            while (node != root && node->children.find(ch) == node->children.end()) {
                node = node->fail;
            }

            if (node->children.find(ch) != node->children.end()) {
                node = node->children[ch];
            } else {
                node = root;
            }

            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;
}