Cod sursa(job #2642939)

Utilizator IulianOleniucIulian Oleniuc IulianOleniuc Data 17 august 2020 19:04:02
Problema Aho-Corasick Scor 25
Compilator cpp-64 Status done
Runda Arhiva educationala Marime 2.67 kb
#include <bits/stdc++.h>
using namespace std;

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

class SuffixAutomaton {
  private:
    struct Node {
        int len, link;
        bool clone = false;
        int next[26] = {};
    };

    int last = 0;
    vector<Node> dag;
    vector<int> dp;

  public:
    SuffixAutomaton(const string& str = "") {
        dag.reserve(str.size() * 2);
        dag.emplace_back();
        dag[0].len = 0;
        dag[0].link = -1;
        last = 0;
        for (char chr : str)
            extend(chr);
    }

    void extend(char chr) {
        int now = dag.size();
        dag.emplace_back();
        dag[now].len = dag[last].len + 1;
        int p = last;
        while (p != -1 && !dag[p].next[chr - 'a']) {
            dag[p].next[chr - 'a'] = now;
            p = dag[p].link;
        }
        if (p == -1)
            dag[now].link = 0;
        else {
            int q = dag[p].next[chr - 'a'];
            if (dag[q].len == dag[p].len + 1)
                dag[now].link = q;
            else {
                int clone = dag.size();
                dag.emplace_back();
                dag[clone].clone = true;
                dag[clone].len = dag[p].len + 1;
                memcpy(dag[clone].next, dag[q].next, sizeof(dag[0].next));
                dag[clone].link = dag[q].link;
                while (p != -1 && dag[p].next[chr - 'a'] == q) {
                    dag[p].next[chr - 'a'] = clone;
                    p = dag[p].link;
                }
                dag[q].link = dag[now].link = clone;
            }
        }
        last = now;
    }

    void precalcDP() {
        dp.resize(dag.size());
        vector<int> nodes(dag.size());
        for (int i = 0; i < (int) dag.size(); i++)
            nodes[i] = i;
        sort(nodes.begin(), nodes.end(), [&](int x, int y) {
            return dag[x].len > dag[y].len;
        });
        for (int node : nodes)
            if (!dag[node].clone)
                dp[node] = 1;
        for (int node : nodes)
            if (dag[node].link != -1)
                dp[dag[node].link] += dp[node];
    }

    int count(const string& str) {
        int node = 0;
        for (char chr : str) {
            if (!dag[node].next[chr - 'a'])
                return 0;
            node = dag[node].next[chr - 'a'];
        }
        return dp[node];
    }
};

int main() {
    string str; fin >> str;
    SuffixAutomaton sa(str);
    sa.precalcDP();

    int q; fin >> q;
    for (int i = 0; i < q; i++) {
        string word; fin >> word;
        fout << sa.count(word) << '\n';
    }
    fout.close();
    return 0;
}