Cod sursa(job #2738893)

Utilizator BAlexandruBorgovan Alexandru BAlexandru Data 6 aprilie 2021 14:55:13
Problema Aho-Corasick Scor 60
Compilator cpp-64 Status done
Runda Arhiva educationala Marime 3.71 kb
#include <fstream>
#include <vector>
#include <queue>
#include <unordered_map>

using namespace std;

ifstream f("ahocorasick.in");
ofstream g("ahocorasick.out");

const int SIGMA = 26;

class trie
{
    private:
        struct trieNode
        {
            int cnt;

            vector < int > words;

            trieNode* suffixLink;
            unordered_map < int, trieNode* > children;

            trieNode();
        };

        trieNode* root;

        int numberOfWords;

        void clearTrie(trieNode*& node);

    public:
        trie();
        ~trie();
        void insertKey(const string& key);
        void computeSuffixLinks();
        void query(const string& text);
};

trie :: trieNode :: trieNode()
{
    cnt = 0;

    suffixLink = NULL;
}

void trie :: clearTrie(trieNode*& node)
{
    for (auto &it : node -> children)
        clearTrie(it.second);

    delete node;
}

trie :: trie()
{
    numberOfWords = 0;
    root = new trieNode;
}

trie :: ~trie()
{
    clearTrie(root);
}

void trie :: insertKey(const string &key)
{
    trieNode* node = root;

    for (int i = 0; i < key.size(); i++)
    {
        int index = key[i] - 'a';

        if (!node -> children[index])
            node -> children[index] = new trieNode;

        node = node -> children[index];
    }

    numberOfWords++;
    node -> words.push_back(numberOfWords);
}

void trie :: computeSuffixLinks()
{
    queue < trieNode* > q;

    for (auto &it : root -> children)
    {
        it.second -> suffixLink = root;
        q.push(it.second);
    }

    while (!q.empty())
    {
        trieNode* node = q.front();
        q.pop();

        for (auto &it : node -> children)
        {
            trieNode* aux = node -> suffixLink;

            while (aux && !aux -> children.count(it.first))
                aux = aux -> suffixLink;

            if (aux && aux -> children.count(it.first))
                it.second -> suffixLink = aux -> children[it.first];
            else
                it.second -> suffixLink = root;

            q.push(it.second);
        }
    }
}

void trie :: query(const string &text)
{
    trieNode* node = root;

    for (int i = 0; i < text.size(); i++)
    {
        int index = text[i] - 'a';

        if (node -> children.count(index))
            node = node -> children[index];
        else
        {
            trieNode* aux = node -> suffixLink;

            while (aux && !aux -> children.count(index))
                aux = aux -> suffixLink;

            if (aux && aux -> children.count(index))
                node = aux -> children[index];
            else
                node = root;
        }

        node -> cnt++;
    }

    queue < trieNode* > q;
    vector < trieNode* > antibfs;

    q.push(root);

    while (!q.empty())
    {
        trieNode* node = q.front();
        q.pop();

        antibfs.push_back(node);

        for (auto &it : node -> children)
            q.push(it.second);
    }

    int* ans = new int[numberOfWords + 1];

    for (int i = antibfs.size() - 1; i >= 0; i--)
    {
        trieNode* node = antibfs[i];

        for (auto word : node -> words)
            ans[word] = node -> cnt;

        if (node -> suffixLink)
            node -> suffixLink -> cnt += node -> cnt;
    }

    for (int i = 1; i <= numberOfWords; i++)
        g << ans[i] << "\n";

    delete[] ans;
}

string text;

trie T;

int main()
{
    f >> text;

    int n; f >> n;

    for (int i = 1; i <= n; i++)
    {
        string word; f >> word;
        T.insertKey(word);
    }

    T.computeSuffixLinks();

    T.query(text);

    return 0;
}