Cod sursa(job #2341931)

Utilizator preda.andreiPreda Andrei preda.andrei Data 12 februarie 2019 13:21:39
Problema Aho-Corasick Scor 65
Compilator cpp-64 Status done
Runda Arhiva educationala Marime 2.9 kb
#include <cstdio>
#include <map>
#include <queue>
#include <string>
#include <vector>

using namespace std;

struct TrieNode
{
    map<char, TrieNode*> sons;
    TrieNode *suffix = nullptr;

    vector<int> indexes;
    int count = 0;

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

bool TrieNode::HasSon(char ch) const
{
    return sons.count(ch) > 0;
}

TrieNode* TrieNode::Son(char ch)
{
    auto it = sons.find(ch);
    if (it != sons.end()) {
        return it->second;
    }

    sons.insert({ch, new TrieNode});
    return sons[ch];
}

void Insert(TrieNode *root, int index, char const *str)
{
    if (*str == '\0') {
        root->indexes.push_back(index);
        return;
    }
    Insert(root->Son(*str), index, str + 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 (const auto &p : node->sons) {
            p.second->suffix = FindSuffixNode(node, p.first);
            order.push_back(p.second);
        }
    }
    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,
                             char const *text)
{
    auto *node = order[0];
    while (*text != '\0') {
        while (!node->HasSon(*text) && node->suffix) {
            node = node->suffix;
        }
        if (node->HasSon(*text)) {
            node = node->Son(*text);
        }
        node->count += 1;
        text += 1;
    }

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

constexpr auto kTextLen = 1000001;
constexpr auto kWordLen = 1001;

char text[kTextLen + 10];
char word[kWordLen + 10];

int main()
{
    FILE *fin = fopen("ahocorasick.in", "r");
    FILE *fout = fopen("ahocorasick.out", "w");

    fscanf(fin, "%s", text);

    int words;
    fscanf(fin, "%d%*c", &words);

    TrieNode *trie = new TrieNode;
    for (auto i = 0; i < words; i += 1) {
        fscanf(fin, "%s", word);
        Insert(trie, i, word);
    }

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

    for (const auto &count : res) {
        fprintf(fout, "%d\n", count);
    }

    return 0;
}