Cod sursa(job #2341944)

Utilizator preda.andreiPreda Andrei preda.andrei Data 12 februarie 2019 13:38:05
Problema Aho-Corasick Scor 95
Compilator cpp-64 Status done
Runda Arhiva educationala Marime 2.94 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;

    TrieNode()
    {
        for (auto i = 0; i < 26; i += 1) {
            sons[i] = nullptr;
        }
    }

    bool HasSon(char ch) const;
    TrieNode* Son(char ch);
};

bool TrieNode::HasSon(char ch) const
{
    return sons[ch - 'a'];
}

TrieNode* TrieNode::Son(char ch)
{
    if (sons[ch - 'a']) {
        return sons[ch - 'a'];
    }
    return sons[ch - 'a'] = new TrieNode;
}

void Insert(TrieNode *root, int index, const string &str, size_t pos = 0)
{
    if (pos >= str.size()) {
        root->indexes.push_back(index);
        return;
    }
    Insert(root->Son(str[pos]), index, str, pos + 1);
}

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

    auto other = father->suffix;
    while (other->suffix && !other->HasSon(ch)) {
        other = other->suffix;
    }

    if (other->HasSon(ch)) {
        other = other->Son(ch);
    }
    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->HasSon(ch)) {
                continue;
            }

            auto son = node->Son(ch);
            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 (const auto &index : node->indexes) {
            count[index] = 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 (const auto &ch : text) {
        while (!node->HasSon(ch) && node->suffix) {
            node = node->suffix;
        }
        if (node->HasSon(ch)) {
            node = node->Son(ch);
        }
        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 (const auto &count : res) {
        fout << count << "\n";
    }

    return 0;
}