Cod sursa(job #2738944)

Utilizator bogdanvladmihaiBogdan Vlad-Mihai bogdanvladmihai Data 6 aprilie 2021 16:33:14
Problema Aho-Corasick Scor 95
Compilator cpp-64 Status done
Runda Arhiva educationala Marime 2.72 kb
#include <bits/stdc++.h>

using namespace std;

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

const int max_s = (int)1e6 + 5;
const int max_n = (int)1e4 + 5;
const int alpha = 26;

struct Trie {
  vector<int> wordIds;
  int cnt;
  Trie *child[alpha];
  Trie *fail;
  Trie() {
    for (int i = 0; i < alpha; i++) {
      child[i] = nullptr;
    }
    fail = nullptr;
    cnt = 0;
    wordIds.clear();
  }
};

int n;

Trie *root = new Trie();

char s[max_s];

char c[max_n];

int sol[max_s];

void insert(Trie *root, char *s, int id) {
  Trie *currentNode = root;
  int n = strlen(s);
  for (int i = 0; i < n; i++) {
    if (currentNode -> child[s[i] - 'a'] == nullptr) {
      currentNode -> child[s[i] - 'a'] = new Trie();
    }
    currentNode = currentNode -> child[s[i] - 'a'];
  }
  currentNode -> wordIds.push_back(id);
}

void aho() {
  queue<Trie *> q;
  root -> fail = nullptr;
  for (int i = 0; i < alpha; i++) {
    if (root -> child[i] != nullptr) {
      q.push(root -> child[i]);
      root -> child[i] -> fail = root;
    }
  }
  while ((int)q.size() > 0) {
    Trie *u = q.front();
    q.pop();
    for (int i = 0; i < alpha; i++) {
      if (u -> child[i] != nullptr) {
        q.push(u -> child[i]);
        Trie *aux = u -> fail;
        while (aux != root && aux -> child[i] == nullptr) {
          aux = aux -> fail;
        }
        if (aux -> child[i] != nullptr) {
          u -> child[i] -> fail = aux -> child[i];
        } else {
          u -> child[i] -> fail = root;
        }
      }
    }
  }
}

void solve() {
  Trie *currentNode = root;
  int n = strlen(s);
  for (int i = 0; i < n; i++) {
    while (currentNode != root && currentNode -> child[s[i] - 'a'] == nullptr) {
      currentNode = currentNode -> fail;
    }
    if (currentNode -> child[s[i] - 'a'] != nullptr) {
      currentNode = currentNode -> child[s[i] - 'a'];
    } else {
      currentNode = root;
    }
    currentNode -> cnt++;
  }
  queue<Trie *> q;
  q.push(root);
  vector<Trie *> antiDfs;
  while ((int)q.size() > 0) {
    Trie *u = q.front();
    q.pop();
    antiDfs.push_back(u);
    for (int i = 0; i < alpha; i++) {
      if (u -> child[i] != nullptr) {
        q.push(u -> child[i]);
      }
    }
  }
  reverse(antiDfs.begin(), antiDfs.end());
  for (Trie *u : antiDfs) {
    for (int v : u -> wordIds) {
      sol[v] += u -> cnt;
    }
    if (u -> fail != nullptr) {
      u -> fail -> cnt += u -> cnt;
    }
  }
}

int main() {
  in >> s;
  in >> n;
  for (int i = 1; i <= n; i++) {
    in >> c;
    insert(root, c, i);
  }
  aho();
  solve();
  for (int i = 1; i <= n; i++) {
    out << sol[i] << "\n";
  }
  return 0;
}