Problem p02589 (C++) — buggy submission, verdict: Time Limit Exceeded

#include <algorithm>
#include <bitset>
#include <cassert>
#include <chrono>
#include <cstring>
#include <fstream>
#include <iostream>
#include <map>
#include <math.h>
#include <numeric>
#include <queue>
#include <set>
#include <stack>
#include <tuple>
#include <unordered_map>
#include <vector>
using namespace std;

const int N = 1e6 + 10, mod = 1e9 + 7;

int trie[N][26];
int Cnt[N][26];
int sz = 1;

char buf[N];

int main() {
  // ios_base::sync_with_stdio(0); cin.tie(0); cout.tie(0);
  int tc = 1;
  // scanf("%d", &tc);
  for (int cn = 1; cn <= tc; cn++) {
    int n;
    scanf("%d", &n);
    vector<string> s(n);
    vector<int> cnt(26, 0);
    for (int i = 0; i < n; i++) {
      scanf("%s", &buf);
      s[i] = buf;
      if (s[i].size() == 1) {
        cnt[s[i][0] - 'a']++;
        continue;
      }
      int u = 0;
      for (int j = (int)s[i].size() - 1; j > 0; j--) {
        int c = s[i][j] - 'a';
        if (!trie[u][c])
          trie[u][c] = sz++;
        u = trie[u][c];
      }
      Cnt[u][s[i][0] - 'a']++;
    }
    long long ans = 0;
    for (int i = 0; i < n; i++) {
      if (s[i].size() == 1)
        continue;
      vector<int> mask(s[i].size(), 0);
      mask[0] = (1 << (s[i][0] - 'a'));
      for (int j = 1; j < s[i].size(); j++)
        mask[j] = (mask[j - 1] | (1 << (s[i][j] - 'a')));
      int u = 0;
      for (int j = s[i].size() - 1; j > 1; j--) {
        u = trie[u][s[i][j] - 'a'];
        for (int c = 0; c < 26; c++) {
          if ((mask[j - 1] >> c) & 1)
            ans += Cnt[u][c];
        }
      }
      for (int c = 0; c < 26; c++) {
        if ((mask[s[i].size() - 1] >> c) & 1)
          ans += cnt[c];
      }
    }
    printf("%lld\\n", ans);
  }
  return 0;
}

Error: TLE

Fixed submission:

#include <algorithm>
#include <bitset>
#include <cassert>
#include <chrono>
#include <cstring>
#include <fstream>
#include <iostream>
#include <map>
#include <math.h>
#include <numeric>
#include <queue>
#include <set>
#include <stack>
#include <tuple>
#include <unordered_map>
#include <vector>
using namespace std;

const int N = 1e6 + 10, mod = 1e9 + 7;

class Hash {
private:
  int n;
  vector<int> pw, invpw, pref;
  int m, base;

public:
  int fp(int a, int b) {
    int ret = 1;
    while (b) {
      if (b & 1)
        ret = 1ll * ret * a % m;
      b >>= 1;
      a = 1ll * a * a % m;
    }
    return ret;
  }

  Hash() {}
  Hash(string &_str, int _base, int _m) {
    n = _str.size();
    base = _base;
    m = _m;
    pw.resize(n + 1, 1);
    invpw.resize(n + 1, 1);
    pref.resize(n + 1, 0);
    pw[1] = base;
    invpw[1] = fp(base, m - 2);
    for (int i = 2; i < n; i++) {
      pw[i] = 1ll * pw[i - 1] * base % m;
      invpw[i] = 1ll * invpw[i - 1] * invpw[1] % m;
    }
    for (int i = 0; i < n; i++) {
      pref[i + 1] = (pref[i] + 1ll * pw[i] * _str[i] % m) % m;
    }
  }

  int get(int i, int len) {
    int ret = (1ll * pref[i + len] - pref[i]) % m;
    ret = 1ll * ret * invpw[i] % m;
    if (ret < 0)
      ret += m;
    return ret;
  }
};

struct hash_pair {
  template <class T1, class T2> size_t operator()(const pair<T1, T2> &p) const {
    auto hash1 = hash<T1>{}(p.first);
    auto hash2 = hash<T2>{}(p.second);
    return hash1 ^ hash2;
  }
};

char buf[N];

int main() {
  // ios_base::sync_with_stdio(0); cin.tie(0); cout.tie(0);
  int tc = 1;
  // scanf("%d", &tc);
  for (int cn = 1; cn <= tc; cn++) {
    int n;
    scanf("%d", &n);
    vector<string> s(n);
    vector<Hash> h1;
    vector<Hash> h2;
    unordered_map<pair<int, int>, int, hash_pair> have[26];
    for (int i = 0; i < n; i++) {
      scanf("%s", &buf);
      s[i] = buf;
      h1.push_back(Hash(s[i], 1337, 1e9 + 123));
      h2.push_back(Hash(s[i], 7331, 1e9 + 321));
      if (s[i].size() == 1) {
        have[s[i][0] - 'a'][{0, 0}]++;
        continue;
      }
      have[s[i][0] - 'a']
          [{h1[i].get(1, s[i].size() - 1), h2[i].get(1, s[i].size() - 1)}]++;
    }
    long long ans = 0;
    for (int i = 0; i < n; i++) {
      if (s[i].size() == 1)
        continue;
      vector<bool> vis(26, 0);
      vis[s[i][0] - 'a'] = 1;
      vis[s[i][1] - 'a'] = 1;
      for (int j = 2; j < s[i].size(); j++) {
        pair<int, int> cur = {h1[i].get(j, s[i].size() - j),
                              h2[i].get(j, s[i].size() - j)};
        for (int c = 0; c < 26; c++) {
          if (!vis[c])
            continue;
          ans += have[c][cur];
        }
        vis[s[i][j] - 'a'] = 1;
      }
      for (int c = 0; c < 26; c++) {
        if (!vis[c])
          continue;
        ans += have[c][{0, 0}];
      }
    }
    printf("%lld\\n", ans);
  }
  return 0;
}