Cod sursa(job #2341969)

Utilizator preda.andreiPreda Andrei preda.andrei Data 12 februarie 2019 14:06:39
Problema Aho-Corasick Scor 95
Compilator cpp-64 Status done
Runda Arhiva educationala Marime 2.75 kb
#include <fstream>
#include <queue>
#include <string>
#include <vector>

using namespace std;

struct TrieNode
{
    TrieNode* sons[26] = {};
    TrieNode *suffix = nullptr;

    vector<int> indexes;
    int count = 0;
};

void Insert(TrieNode *root, int index, const string &str)
{
    for (size_t i = 0; i < str.size(); i += 1) {
        if (!root->sons[str[i] - 'a']) {
            root->sons[str[i] - 'a'] = new TrieNode;
        }
        root = root->sons[str[i] - 'a'];
    }
    root->indexes.push_back(index);
}

TrieNode* FindSuffixNode(TrieNode *father, char ch)
{
    if (!father->suffix) {
        return father;
    }

    auto other = father->suffix;
    while (other->suffix && !other->sons[ch - 'a']) {
        other = other->suffix;
    }

    if (other->sons[ch - 'a']) {
        other = other->sons[ch - 'a'];
    }
    return other;
}

vector<TrieNode*> MakeAutomaton(TrieNode *root)
{
    vector<TrieNode*> order;
    order.push_back(root);

    size_t index = 0;
    while (index < order.size()) {
        auto node = order[index++];

        for (auto ch = 'a'; ch <= 'z'; ch += 1) {
            if (!node->sons[ch - 'a']) {
                continue;
            }

            auto son = node->sons[ch - 'a'];
            son->suffix = FindSuffixNode(node, ch);
            order.push_back(son);
        }
    }
    return order;
}

void ExtractCount(const vector<TrieNode*> &order, vector<int> &count)
{
    for (int i = order.size() - 1; i >= 0; i -= 1) {
        auto node = order[i];
        for (size_t j = 0; j < node->indexes.size(); j += 1) {
            count[node->indexes[j]] = node->count;
        }
        if (node->suffix) {
            node->suffix->count += node->count;
        }
    }
}

vector<int> CountAppearances(int words,
                             const vector<TrieNode*> &order,
                             const string &text)
{
    auto *node = order[0];
    for (size_t i = 0; i < text.size(); i += 1) {
        while (!node->sons[text[i] - 'a'] && node->suffix) {
            node = node->suffix;
        }
        if (node->sons[text[i] - 'a']) {
            node = node->sons[text[i] - 'a'];
        }
        node->count += 1;
    }

    vector<int> count(words, 0);
    ExtractCount(order, count);
    return count;
}

int main()
{
    ifstream fin("ahocorasick.in");
    ofstream fout("ahocorasick.out");

    string text;
    getline(fin, text);

    int words;
    fin >> words;
    fin.get();

    TrieNode *trie = new TrieNode;
    for (auto i = 0; i < words; i += 1) {
        string word;
        getline(fin, word);
        Insert(trie, i, word);
    }

    auto order = MakeAutomaton(trie);
    auto res = CountAppearances(words, order, text);

    for (auto i = 0; i < words; i += 1) {
        fout << res[i] << "\n";
    }

    return 0;
}