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

#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>

using namespace std;
using namespace __gnu_pbds;

typedef long long ll;

const int N = 1000010;
const int MOD = 1e9 + 7;
const ll P[2] = {97, 1000003};

ll bigMod(ll a, ll e) {
  if (e == -1)
    e = MOD - 2;
  ll ret = 1;
  while (e) {
    if (e & 1)
      ret = ret * a % MOD;
    a = a * a % MOD, e >>= 1;
  }
  return ret;
}

int n, id[N], len[N];
char tmp[N];
string a[N];
map<ll, vector<int>> cnt;
ll pwr[2][N], inv[2][N], h[2][N];

inline ll get(int l, int r) {
  if (l > r)
    return 0;
  ll one = (h[0][r] - h[0][l - 1]) * inv[0][l] % MOD;
  ll two = (h[1][r] - h[1][l - 1]) * inv[1][l] % MOD;
  if (one < 0)
    one += MOD;
  if (two < 0)
    two += MOD;
  return (one << 31LL) | two;
}

int main() {
  for (int it = 0; it < 2; ++it) {
    pwr[it][0] = inv[it][0] = 1;
    ll cur = P[it], inv_cur = bigMod(cur, -1);
    for (int i = 1; i < N; ++i) {
      pwr[it][i] = pwr[it][i - 1] * cur % MOD;
      inv[it][i] = inv[it][i - 1] * inv_cur % MOD;
    }
  }
  cin >> n;
  for (int i = 1; i <= n; ++i) {
    id[i] = i;
    scanf("%s", tmp);
    a[i] = tmp, len[i] = strlen(tmp);
  }
  sort(id + 1, id + n + 1, [](int i, int j) { return len[i] < len[j]; });
  ll ans = 0;
  for (int _i = 1; _i <= n; ++_i) {
    int i = id[_i], l = len[i];
    for (int it = 0; it < 2; ++it) {
      h[it][0] = 0;
      for (int j = 1; j <= l; ++j) {
        h[it][j] = (h[it][j - 1] + (a[i][j - 1] - 'a' + 1) * pwr[it][j]) % MOD;
      }
    }
    int prv = 1 << (a[i][0] - 'a');
    for (int j = 1; j <= l; ++j) {
      ll cur = get(j + 1, l);
      if (cnt.find(cur) != cnt.end()) {
        vector<int> vec = cnt[cur];
        for (int k = 0; k < 26; ++k)
          if (prv & 1 << k) {
            ans += vec[k];
          }
      }
      if (j < l)
        prv |= 1 << (a[i][j] - 'a');
    }
    ll cur = get(2, l);
    if (cnt.find(cur) == cnt.end())
      cnt[cur] = vector<int>(26, 0);
    ++cnt[cur][a[i][0] - 'a'];
  }
  cout << ans << '\\n';
  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>

using namespace std;
using namespace __gnu_pbds;

typedef long long ll;

const int N = 1000010;
const int MOD = 1e9 + 7;
const ll P[2] = {97, 1000003};
const ll RANDOM = 928512321984217893LL;

struct chash {
  ll operator()(ll x) const { return x ^ RANDOM; }
};

ll bigMod(ll a, ll e) {
  if (e == -1)
    e = MOD - 2;
  ll ret = 1;
  while (e) {
    if (e & 1)
      ret = ret * a % MOD;
    a = a * a % MOD, e >>= 1;
  }
  return ret;
}

int n, id[N], len[N];
char tmp[N];
string a[N];
ll pwr[2][N], inv[2][N], h[2][N];
gp_hash_table<ll, int, chash> cnt[30];

inline ll get(int l, int r) {
  if (l > r)
    return 0;
  ll one = (h[0][r] - h[0][l - 1]) * inv[0][l] % MOD;
  ll two = (h[1][r] - h[1][l - 1]) * inv[1][l] % MOD;
  if (one < 0)
    one += MOD;
  if (two < 0)
    two += MOD;
  return (one << 31LL) | two;
}

int main() {
  for (int it = 0; it < 2; ++it) {
    pwr[it][0] = inv[it][0] = 1;
    ll cur = P[it], inv_cur = bigMod(cur, -1);
    for (int i = 1; i < N; ++i) {
      pwr[it][i] = pwr[it][i - 1] * cur % MOD;
      inv[it][i] = inv[it][i - 1] * inv_cur % MOD;
    }
  }
  cin >> n;
  for (int i = 1; i <= n; ++i) {
    id[i] = i;
    scanf("%s", tmp);
    a[i] = tmp, len[i] = strlen(tmp);
  }
  sort(id + 1, id + n + 1, [](int i, int j) { return len[i] < len[j]; });
  ll ans = 0;
  for (int _i = 1; _i <= n; ++_i) {
    int i = id[_i], l = len[i];
    for (int it = 0; it < 2; ++it) {
      h[it][0] = 0;
      for (int j = 1; j <= l; ++j) {
        h[it][j] = (h[it][j - 1] + (a[i][j - 1] - 'a' + 1) * pwr[it][j]) % MOD;
      }
    }
    int prv = 1 << (a[i][0] - 'a');
    for (int j = 1; j <= l; ++j) {
      ll cur = get(j + 1, l);
      for (int k = 0; k < 26; ++k)
        if (prv & 1 << k) {
          ans += cnt[k][cur];
        }
      if (j < l)
        prv |= 1 << (a[i][j] - 'a');
    }
    ll cur = get(2, l);
    ++cnt[a[i][0] - 'a'][cur];
  }
  cout << ans << '\\n';
  return 0;
}