Cod sursa(job #2002696)

Utilizator borcanirobertBorcani Robert borcanirobert Data 20 iulie 2017 16:54:26
Problema Aho-Corasick Scor 10
Compilator cpp Status done
Runda Arhiva educationala Marime 3.45 kb
#include <iostream>
#include <fstream>
#include <vector>
using namespace std;

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

class Trie{
public:
    static string s;

    Trie();
    void AddWord( int pos );   // only after adding all the words I must make blue and green edges
    void MakeBlueEdges();
    void MakeGreenEdges();
    void _MakeGreenEdges();
    void Solve( int pos );
private:
    int word;    // is this the end of a word from the dictionary?
    Trie *sons[26];
    Trie *blueEdge, *greenEdge;
    Trie *parent;
    char letter;
};
Trie *root = new Trie();
string Trie::s = "";
string text;
int N, M;
vector< pair<int, int> > sol;
vector<int> cnt;
int nr;

int main()
{
    fin >> text;
    fin >> M;
    cnt = vector<int>(M + 1);

    for ( int i = 0; i < M; ++i )
    {
        ++nr;
        fin >> Trie::s;
        N = Trie::s.size();
        root->AddWord(0);
    }

    root->MakeBlueEdges();
    root->MakeGreenEdges();

    Trie::s = text;
    N = text.size();
    root->Solve(0);

    for ( int i = 1; i <= M; ++i )
        fout << cnt[i] << '\n';

    fin.close();
    fout.close();
    return 0;
}

Trie::Trie()
{
    word = false;
    for ( int i = 0; i < 26; ++i )
        sons[i] = nullptr;
    blueEdge = greenEdge = parent = nullptr;
    letter = 'a';
}

void Trie::AddWord( int p )
{
    if ( p == Trie::s.size() )
    {
        word = nr;
        return;
    }

    int pos = s[p] - 'a';
    if ( sons[pos] == nullptr )
    {
        sons[pos] = new Trie();
        sons[pos]->letter = s[p];
        sons[pos]->parent = this;
    }
    sons[pos]->AddWord(p + 1);
}

void Trie::MakeBlueEdges()
{
    if ( this != root )
    {
        int p = letter - 'a';
        auto _parent = parent;
        while ( _parent != root && _parent->blueEdge->sons[p] == nullptr )
            _parent = _parent->parent;
        blueEdge = _parent != root && _parent->blueEdge->sons[p] != nullptr
                && _parent->blueEdge->sons[p] != this ? _parent->blueEdge->sons[p] : root;
    }

    for ( int i = 0; i < 26; ++i )
        if ( sons[i] != nullptr )
            sons[i]->MakeBlueEdges();
}

void Trie::MakeGreenEdges()
{
    for ( int i = 0; i < 26; ++i )
        if ( sons[i] != nullptr )
            sons[i]->_MakeGreenEdges();
}

void Trie::_MakeGreenEdges()
{
    if ( this == root || greenEdge != nullptr )
        return;

    if ( blueEdge->word )
        greenEdge = blueEdge;
    else
    {
        blueEdge->_MakeGreenEdges();
        if ( blueEdge != root && blueEdge->greenEdge != nullptr )
            greenEdge = blueEdge->greenEdge;
    }

    for ( int i = 0; i < 26; ++i )
        if ( sons[i] != nullptr )
            sons[i]->_MakeGreenEdges();
}

void Trie::Solve( int pos )
{
    if ( pos > N )
        return;

    if ( this == root )
    {
        int l = Trie::s[pos] - 'a';
        if ( sons[l] == nullptr )
            this->Solve(pos + 1);
        else
            sons[l]->Solve(pos + 1);
    }
    else
    {
        if ( word )
            ++cnt[word];
        auto greens = greenEdge;
        while ( greens != nullptr )
            ++cnt[greens->word], greens = greens->greenEdge;

        if ( pos < N )
        {
            int l = Trie::s[pos] - 'a';
            auto curr = this;
            while ( curr != root && curr->sons[l] == nullptr )
                curr = curr->blueEdge;
            if ( curr->sons[l] != nullptr )
                curr->sons[l]->Solve(pos + 1);
            else
                root->Solve(pos + 1);
        }
    }
}