Problem p03974 (C++)

Buggy submission (status: Time Limit Exceeded):
// VSCF.cpp : Defines the entry point for the console application.
//
#include <algorithm>
#include <cassert>
#include <cmath>
#include <cstdio>
#include <deque>
#include <functional>
#include <iomanip>
#include <iostream>
#include <iterator>
#include <list>
#include <map>
#include <set>
#include <sstream>
#include <string>
#include <vector>
using namespace std;
typedef long long LL;
typedef pair<int, int> PII;
#define MP make_pair
#define FOR(v, p, k) for (int v = p; v <= k; ++v)
#define FORD(v, p, k) for (int v = p; v >= k; --v)
#define REP(i, n) for (int i = 0; i < (n); ++i)
#define VAR(v, i) __typeof(i) v = (i)
#define FOREACH(i, c) for (VAR(i, (c).begin()); i != (c).end(); ++i)
#define PB push_back
#define ST first
#define ND second
#define SIZE(x) (int)x.size()
#define ALL(c) c.begin(), c.end()
#define int long long
#define MAXN 1000010
typedef long double LD;

vector<string> inp2;

struct Trie {
  Trie *sons[26];
  Trie *father;
  int fatherPos;
  int isWord;
  int sonsCo = 0;
  int size;
  static int getSize(Trie *t) {
    if (t) {
      return t->size;
    } else {
      return 0;
    }
  }
  Trie() {
    REP(i, 26) { sons[i] = 0; }
    father = 0;
    isWord = 0;
    size = 0;
  }
  Trie *addWord(int idx, const string &s, int no) {
    if (idx == s.size()) {
      isWord = no;
      return this;
    }
    int letter = s[idx] - 'a';
    if (sons[letter]) {
      sons[letter]->addWord(idx + 1, s, no);
    } else {
      sons[letter] = new Trie();
      sonsCo++;
      sons[letter]->addWord(idx + 1, s, no);
    }
  }
  Trie *compress() {
    if (!isWord && sonsCo == 1) {
      int i = 0;
      while (!sons[i]) {
        i++;
      }
      return sons[i];
    }
    if (sonsCo) {
      REP(i, 26) {
        if (sons[i]) {
          sons[i] = sons[i]->compress();
        }
      }
    }
    return this;
  }
  int findSize() {
    size = 0;
    if (isWord) {
      size++;
    }
    REP(i, 26) {
      if (sons[i]) {
        size += sons[i]->findSize();
      }
    }
    REP(i, 26) {
      if (sons[i]) {
        sons[i]->father = this;
        sons[i]->fatherPos = i;
      }
    }
    return size;
  }
  int findWord(int idx, const string &s) {
    if (idx == 1 && isWord) {
      return isWord;
    }
    if (isWord) {
      idx--;
    }
    int i = 0;
    while (getSize(sons[s[i] - 'a']) < idx) {
      idx -= getSize(sons[s[i] - 'a']);
      i++;
    }
    return sons[s[i] - 'a']->findWord(idx, s);
  }
  int findBefore(int sonIdx, const string &s) {
    int i = 0;
    int res = 0;
    if (isWord) {
      res++;
    }
    while (s[i] - 'a' != sonIdx) {
      res += getSize(sons[s[i] - 'a']);
      i++;
    }
    if (father) {
      res += father->findBefore(fatherPos, s);
    }
    // cerr << size << " " << getSize(sons[0]) << " " << getSize(sons[1]) << "
    // "<< res << endl;
    return res;
  }
  int findPos(const string &s) {
    int res = 1;
    if (father) {
      res += father->findBefore(fatherPos, s);
    }
    return res;
  }
};

#undef int
int main() {
#define int long long
  int n;
  cin >> n;
  Trie *t = new Trie();
  vector<Trie *> inp(n);
  REP(i, n) {
    string s;
    cin >> s;
    inp[i] = t->addWord(0, s, i + 1);
    inp2.push_back(s);
  }
  t = t->compress();
  t->findSize();
  int q;
  cin >> q;
  REP(i, q) {
    int k;
    string s;
    cin >> k >> s;
    cout << inp[k - 1]->findPos(s) << "\n";
  }
  return 0;
}


Fixed submission:
// VSCF.cpp : Defines the entry point for the console application.
//
#include <algorithm>
#include <cassert>
#include <cmath>
#include <cstdio>
#include <deque>
#include <functional>
#include <iomanip>
#include <iostream>
#include <iterator>
#include <list>
#include <map>
#include <set>
#include <sstream>
#include <string>
#include <vector>
using namespace std;
typedef long long LL;
typedef pair<int, int> PII;
#define MP make_pair
#define FOR(v, p, k) for (int v = p; v <= k; ++v)
#define FORD(v, p, k) for (int v = p; v >= k; --v)
#define REP(i, n) for (int i = 0; i < (n); ++i)
#define VAR(v, i) __typeof(i) v = (i)
#define FOREACH(i, c) for (VAR(i, (c).begin()); i != (c).end(); ++i)
#define PB push_back
#define ST first
#define ND second
#define SIZE(x) (int)x.size()
#define ALL(c) c.begin(), c.end()
#define int long long
#define MAXN 1000010
typedef long double LD;

vector<string> inp2;

struct Trie {
  Trie *sons[26];
  Trie *father;
  int fatherPos;
  int isWord;
  int sonsCo = 0;
  int size;
  static int getSize(Trie *t) {
    if (t) {
      return t->size;
    } else {
      return 0;
    }
  }
  Trie() {
    REP(i, 26) { sons[i] = 0; }
    father = 0;
    isWord = 0;
    size = 0;
  }
  Trie *addWord(int idx, const string &s, int no) {
    if (idx == s.size()) {
      isWord = no;
      return this;
    }
    int letter = s[idx] - 'a';
    if (sons[letter]) {
      sons[letter]->addWord(idx + 1, s, no);
    } else {
      sons[letter] = new Trie();
      sonsCo++;
      sons[letter]->addWord(idx + 1, s, no);
    }
  }
  Trie *compress() {
    if (!isWord && sonsCo == 1) {
      int i = 0;
      while (!sons[i]) {
        i++;
      }
      return sons[i]->compress();
    }
    if (sonsCo) {
      REP(i, 26) {
        if (sons[i]) {
          sons[i] = sons[i]->compress();
        }
      }
    }
    return this;
  }
  int findSize() {
    size = 0;
    if (isWord) {
      size++;
    }
    REP(i, 26) {
      if (sons[i]) {
        size += sons[i]->findSize();
      }
    }
    REP(i, 26) {
      if (sons[i]) {
        sons[i]->father = this;
        sons[i]->fatherPos = i;
      }
    }
    return size;
  }
  int findWord(int idx, const string &s) {
    if (idx == 1 && isWord) {
      return isWord;
    }
    if (isWord) {
      idx--;
    }
    int i = 0;
    while (getSize(sons[s[i] - 'a']) < idx) {
      idx -= getSize(sons[s[i] - 'a']);
      i++;
    }
    return sons[s[i] - 'a']->findWord(idx, s);
  }
  int findBefore(int sonIdx, const string &s) {
    int i = 0;
    int res = 0;
    if (isWord) {
      res++;
    }
    while (s[i] - 'a' != sonIdx) {
      res += getSize(sons[s[i] - 'a']);
      i++;
    }
    if (father) {
      res += father->findBefore(fatherPos, s);
    }
    // cerr << size << " " << getSize(sons[0]) << " " << getSize(sons[1]) << "
    // "<< res << endl;
    return res;
  }
  int findPos(const string &s) {
    int res = 1;
    if (father) {
      res += father->findBefore(fatherPos, s);
    }
    return res;
  }
};

#undef int
int main() {
#define int long long
  int n;
  cin >> n;
  Trie *t = new Trie();
  vector<Trie *> inp(n);
  REP(i, n) {
    string s;
    cin >> s;
    inp[i] = t->addWord(0, s, i + 1);
    inp2.push_back(s);
  }
  t = t->compress();
  t->findSize();
  int q;
  cin >> q;
  REP(i, q) {
    int k;
    string s;
    cin >> k >> s;
    cout << inp[k - 1]->findPos(s) << "\n";
  }
  return 0;
}