Cod sursa(job #2898098)

Utilizator AlexNeaguAlexandru AlexNeagu Data 5 mai 2022 23:49:37
Problema Aho-Corasick Scor 70
Compilator cpp-64 Status done
Runda Arhiva educationala Marime 2.03 kb
#include<bits/stdc++.h>
//#define int long long
using namespace std;
typedef long long ll;
ifstream in("ahocorasick.in");
ofstream out("ahocorasick.out");
struct AhoCorasick {
  struct Node {
    int link;
    map<char, int> leg;
  };
  vector<Node> T;
  vector<int> leafs, allNodes;
  int root = 0, nodes = 1, cntNodes;

  AhoCorasick(int sz) : T(sz), cntNodes(sz) {}

  // Adds a word to trie and returns the end node
  void AddWord(const string &word) {
    int node = root;
    for (auto c : word) {
      auto &nxt = T[node].leg[c];
      if (nxt == 0) nxt = nodes++;
      node = nxt;
    }
    leafs.push_back(node);
  }

  // Advances from a node with a character (like an automaton)
  int Advance(int node, char chr) {
    while (node != -1 && T[node].leg.count(chr) == 0)
      node = T[node].link;
    if (node == -1) return root;
    return T[node].leg[chr];
  }

  // Builds links
  void BuildLinks() {
    queue<int> Q;
    Q.push(root);
    allNodes.push_back(root);
    T[root].link = -1;

    while (!Q.empty()) {
      int node = Q.front();
      Q.pop();
      
      for (auto &p : T[node].leg) {
        int vec = p.second;
        char chr = p.first;
        T[vec].link = Advance(T[node].link, chr);
        Q.push(vec);
        allNodes.push_back(vec);
      }
    }
  }
  void ask(string &s) {
	vector<int> freq(cntNodes);
	int node = root;
	for(auto it : s) {
		while(node != root && T[node].leg.count(it) == 0) {
			node = T[node].link;
		}
		if(T[node].leg.count(it)) {
			node = T[node].leg[it];
		}
		freq[node]++;
	}
	for(int i = (int) allNodes.size() - 1; i > 0; --i) {
		int nod = allNodes[i];
		freq[T[nod].link] += freq[nod];
	}
	for(auto it : leafs) {
		out << freq[it] << '\n';
	}
  }
};
int32_t main() {
	ios_base::sync_with_stdio(false);
	cin.tie(0);
	string s;
	in >> s;
	int n;
	in >> n;
	int totalSize = 0;
	vector<string> input(n);
	for(int i = 0; i < n; ++i) {
		in >> input[i];
		totalSize += (int) input[i].size();
	}
	AhoCorasick A(totalSize * 2);
	for(auto it : input) {
		A.AddWord(it);
	}
	A.BuildLinks();
	A.ask(s);
	return 0;
}