Cod sursa(job #3156819)

Utilizator AleXutzZuDavid Alex Robert AleXutzZu Data 13 octombrie 2023 12:25:19
Problema Aho-Corasick Scor 90
Compilator cpp-64 Status done
Runda Arhiva educationala Marime 2.78 kb
#include <iostream>
#include <vector>
#include <queue>
#include <fstream>
#include <unordered_map>

std::vector<int> ans(100, 0);

struct Trie {
private:
    struct Vertex {
        int next[26] = {0};
        int fail = -1;

        Vertex() {
            for (int &i: next) i = -1;
        }

        int output = -1;
    };

    std::vector<Vertex> cont;

    int &next(int node, char c) {
        return cont[node].next[c - 'a'];
    }

    int &output(int node) {
        return cont[node].output;
    }

    int &fail(int node) {
        return cont[node].fail;
    }

    void push_links() {
        std::queue<int> queue;
        for (char i = 'a'; i <= 'z'; ++i) {
            int state = next(0, i);
            if (state > 0) {
                queue.push(next(0, i));
                fail(state) = 0;
            }
        }

        while (!queue.empty()) {
            int top = queue.front();
            queue.pop();

            for (char i = 'a'; i <= 'z'; ++i) {
                int &current = next(top, i);
                if (current == -1) continue;
                queue.push(current);
                int state = fail(top);
                while (next(state, i) == -1) state = fail(state);
                fail(current) = next(state, i);
            }
        }
    }

public:

    Trie() {
        cont.emplace_back();
    }

    void add_word(const std::string &str, int idx) {
        int node = 0;

        for (char i: str) {
            if (next(node, i) == -1) {
                next(node, i) = (int) cont.size();
                cont.emplace_back();
            }
            node = next(node, i);
        }
        output(node) = idx;
    }

    void compute() {
        fail(0) = 0;
        for (char i = 'a'; i <= 'z'; ++i) {
            if (next(0, i) == -1) next(0, i) = 0;
        }

        push_links();
    }

    void advance(int &state, char c) {
        while (next(state, c) == -1) state = fail(state);

        state = next(state, c);
        if (output(state) != -1) ans[output(state)]++;


        int cpy = fail(state);
        while (cpy != 0) {
            if (output(cpy) != -1) ans[output(cpy)]++;
            cpy = fail(cpy);
        }
    }
};

int main() {
    std::ifstream input("ahocorasick.in");
    std::ofstream output("ahocorasick.out");
    Trie trie;
    std::string s;
    int n;
    input >> s >> n;

    std::unordered_map<std::string, int> map;

    std::vector<std::string> words(n);
    for (int i = 0; i < n; ++i) {
        input >> words[i];

        if (map[words[i]] == 0) map[words[i]] = i + 1;
    }

    for (const auto &[str, idx]: map) {
        trie.add_word(str, idx);
    }

    int state = 0;
    trie.compute();
    for (char i: s) {
        trie.advance(state, i);
    }

    for (int i = 0; i < n; ++i) output << ans[map[words[i]]] << '\n';
    return 0;
}