Cod sursa(job #1490735)

Utilizator salam1Florin Salam salam1 Data 24 septembrie 2015 02:47:16
Problema Aho-Corasick Scor 100
Compilator cpp Status done
Runda Arhiva educationala Marime 2.5 kb
#include <cstdio>
#include <vector>
#include <queue>
#include <cstring>
#include <algorithm>
using namespace std;
const int GAMMA = 26;
const int NMAX = 1000505;
const int LMAX = 10505;
const int HMAX = 105;

struct trieNode {
  trieNode* children[GAMMA];
  trieNode* nxtBest;
  trieNode* par;
  int cnt, whichChild;
  vector<int> end;

  trieNode() {
    cnt = 0;
    memset(children, 0, sizeof(children));
    par = NULL;
  }

  trieNode(trieNode* _par, int _whichChild) {
    par = _par;
    whichChild = _whichChild;
    cnt = 0;
    memset(children, 0, sizeof(children));
  }
};
trieNode* root = new trieNode();
int n, sol[HMAX];
char A[NMAX], word[NMAX];
vector<trieNode*> Q;

void insertTrie(trieNode* node, char *s, int idx) {
  if (*s == '\n') {
    (node -> end).push_back(idx);
    return ;
  }

  if (node -> children[*s - 'a'] == NULL) {
    node -> children[*s - 'a'] = new trieNode(node, *s - 'a');
  }
  insertTrie(node -> children[*s - 'a'], s + 1, idx); 
}

void read() {
  fgets(A + 1, NMAX, stdin);
  scanf("%d\n", &n);
  for (int i = 1; i <= n; i++) {
    fgets(word + 1, NMAX, stdin);
    insertTrie(root, word + 1, i);
  }
}

trieNode* getBestMatch(trieNode* node, int idx) {
  trieNode* best = node;
  while (best != root && best -> children[idx] == NULL) {
    best = best -> nxtBest;
  }
  if (best -> children[idx] != NULL)
    best = best -> children[idx];

  return best;
}

void prepare() {
  // topological sort per level using a queue
  Q.push_back(root);
  for (size_t i = 0; i < Q.size(); i++) {
    trieNode* node = Q[i];
    for (int j = 0; j < GAMMA; j++) {
      if (node -> children[j] != NULL)
        Q.push_back(node -> children[j]);
    }
  }

  // infer fail edge
  for (auto& node: Q) {
    node -> nxtBest = root;
    if (node != root) {
      node -> nxtBest = getBestMatch((node -> par) -> nxtBest, node -> whichChild);
      if (node -> nxtBest == node) {
        node -> nxtBest = root;
      }
    }
  }

  int m = strlen(A + 1) - 1;
  trieNode* lcp = root;
  for (int i = 1; i <= m; i++) {
    lcp = getBestMatch(lcp, A[i] - 'a');
    if (lcp != root) {
      lcp -> cnt++;
    }
  }

  // update count
  reverse(Q.begin(), Q.end());
  for (auto& node: Q) {
    (node -> nxtBest) -> cnt += node -> cnt;
    for (auto x: (node -> end)) {
      sol[x] = node -> cnt;
    }
  }
}

void printSol() {
  for (int i = 1; i <= n; i++)
    printf("%d\n", sol[i]);
}

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

  read();
  prepare();
  printSol();
  return 0;
}