Cod sursa(job #2541941)

Utilizator rapunzelMihnea Andreescu rapunzel Data 9 februarie 2020 10:50:43
Problema Aho-Corasick Scor 95
Compilator cpp-64 Status done
Runda Arhiva educationala Marime 1.91 kb
#include <iostream>
#include <cstdio>
#include <vector>
#include <queue>
#include <algorithm>

using namespace std;

struct trie {
  int sol;
  int len;
  vector<int> ids;
  trie *kids[26];
  trie *go;

  trie() {
    sol = 0;
    len = 0;
    ids.clear();
    go = 0;
    for (int x = 0; x < 26; x++) {
      kids[x] = 0;
    }
  }

};

trie *root = new trie;
vector<int> ans;

void ins(int i, string s) {
  trie *cur = root;
  for (auto &c : s) {
    int x = c - 'a';
    if (!cur->kids[x]) {
      cur->kids[x] = new trie;
      cur->kids[x]->len = cur->len + 1;
    }
    cur = cur->kids[x];
  }
  cur->ids.push_back(i);
}

int main() {
  freopen ("ahocorasick.in", "r", stdin);
  freopen ("ahocorasick.out", "w", stdout);

  string s;
  int cnt;
  cin >> s >> cnt;
  for (int i = 0; i < cnt; i++) {
    string t;
    cin >> t;
    ins(i, t);
    ans.push_back(0);
  }

  root->go = root;
  vector<trie*> rev_ord;
  queue<trie*> q;
  q.push(root);

  while (!q.empty()) {
    trie *cur = q.front();
    rev_ord.push_back(cur);
    q.pop();
    for (int x = 0; x < 26; x++) {
      if (cur->kids[x]) {
        trie *aux = cur->go;
        while (aux != root && !aux->kids[x]) {
          aux = aux->go;
        }
        if (cur->len == 0 || !aux->kids[x]) {
          cur->kids[x]->go = root;
        } else {
          cur->kids[x]->go = aux->kids[x];
        }
        q.push(cur->kids[x]);
      }
    }
  }
  reverse(rev_ord.begin(), rev_ord.end());

  trie *cur = root;
  for (auto &c : s) {
    int x = c - 'a';
    while (cur != root && !cur->kids[x]) {
      cur = cur->go;
    }
    if (cur->kids[x]) {
      cur = cur->kids[x];
    }
    cur->sol++;
  }

  for (auto &it : rev_ord) {
    it->go->sol += it->sol;
    for (auto &i : it->ids) {
      ans[i] += it->sol;
    }
  }

  for (auto &x : ans) {
    cout << x << "\n";
  }

}