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

// #pragma GCC optimize("Ofast")

#ifdef LOCAL
#include "debug.h"
#else
#include <bits/stdc++.h>
#define debug(x...)
#endif

// #define int ll

using namespace std;
typedef long long ll;
typedef long double ld;
typedef pair<int, int> pii;
#define sz(x) int((x).size())

#ifdef ONLINE_JUDGE
mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());
#else
mt19937 rng(1000 - 7);
#endif

const int N = 2e5 + 10;
const int M = 1e6 + 10;
const int MOD = 998244353;
// const int MOD = 1e9 + 7;
const ld eps = 1e-6;
const pii dir[] = {{0, 1}, {0, -1}, {1, 0}, {-1, 0}};

const int K = 2;
const ll primes[] = {227, 229};
ll pw[M][K];

struct custom_hash {
  static uint64_t splitmix64(uint64_t x) {
    x += 0x9e3779b97f4a7c15;
    x = (x ^ (x >> 30)) * 0xbf58476d1ce4e5b9;
    x = (x ^ (x >> 27)) * 0x94d049bb133111eb;
    return x ^ (x >> 31);
  }

  inline size_t operator()(pair<ll, ll> x) const {
    return splitmix64(x.first ^ x.second);
  }
};

int n;
vector<string> v[M], a;
unordered_map<pair<ll, ll>, int, custom_hash> t[26];

signed main() {
#ifdef LOCAL
  freopen("input.txt", "r", stdin);
  freopen("output.txt", "w", stdout);
#endif
  cout << fixed << setprecision(9);
  ios::sync_with_stdio(false), cin.tie(0), cout.tie(0);

  cin >> n;
  a.reserve(n);

  for (int i = 0; i < n; i++) {
    string s;
    cin >> s;

    v[sz(s)].push_back(s);
  }

  for (int i = 0; i < M; i++) {
    for (string &s : v[i]) {
      a.push_back(s);
    }
  }

  for (int i = 0; i < K; i++) {
    pw[0][i] = 1;
  }

  for (int i = 1; i < M; i++) {
    for (int j = 0; j < K; j++) {
      pw[i][j] = (pw[i - 1][j] * primes[j]) % MOD;
    }
  }

  for (int i = 0; i < 26; i++) {
    t[i].max_load_factor(0.25);
    t[i].rehash(M);
  }

  ll ans = 0;
  for (string s : a) {
    debug(s);

    int p[26] = {0};
    for (char c : s) {
      p[c - 'a']++;
    }

    for (int j = 0; j < 26; j++) {
      if (p[j]) {
        ans += t[j][make_pair(0LL, 0LL)];
      }
    }

    ll h[K] = {0};
    for (int i = sz(s) - 1; i > 0; i--) {
      for (int j = 0; j < K; j++) {
        h[j] = (h[j] + (s[i] - 'a' + 1) * pw[sz(s) - i - 1][j]) % MOD;
      }

      p[s[i] - 'a']--;
      auto curH = make_pair(h[0], h[1]);
      for (int j = 0; j < 26; j++) {
        if (p[j] && t[j].count(curH)) {
          ans += t[j][curH];
        }
      }
    }

    t[s.front() - 'a'][make_pair(h[0], h[1])]++;
  }

  cout << ans << endl;

  return 0;
}


Error: TLE

Fixed submission:

#ifdef LOCAL
#include "debug.h"
#else
#include <bits/stdc++.h>
#define debug(x...)
#endif

// #define int ll
// #pragma GCC optimize("Ofast")

using namespace std;
typedef long long ll;
typedef long double ld;
typedef pair<int, int> pii;
#define sz(x) int((x).size())

#ifdef ONLINE_JUDGE
mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());
#else
mt19937 rng(1000 - 7);
#endif

const int N = 2e5 + 10;
const int M = 1e6 + 10;
const int MOD = 998244353;
// const int MOD = 1e9 + 7;
const ld eps = 1e-6;
const pii dir[] = {{0, 1}, {0, -1}, {1, 0}, {-1, 0}};

const int K = 2;
const ll primes[] = {227, 229};
ll pw[M][K];

struct custom_hash {
  static uint64_t splitmix64(uint64_t x) {
    x += 0x9e3779b97f4a7c15;
    x = (x ^ (x >> 30)) * 0xbf58476d1ce4e5b9;
    x = (x ^ (x >> 27)) * 0x94d049bb133111eb;
    return x ^ (x >> 31);
  }

  inline size_t operator()(pair<ll, ll> x) const {
    return splitmix64(x.first ^ x.second);
  }
};

int n;
vector<string> v[M], a;
unordered_map<pair<ll, ll>, int, custom_hash> t[26];

signed main() {
#ifdef LOCAL
  freopen("input.txt", "r", stdin);
  freopen("output.txt", "w", stdout);
#endif
  cout << fixed << setprecision(9);
  ios::sync_with_stdio(false), cin.tie(0), cout.tie(0);

  cin >> n;
  a.reserve(n);

  for (int i = 0; i < n; i++) {
    string s;
    cin >> s;

    v[sz(s)].push_back(s);
  }

  for (int i = 0; i < M; i++) {
    for (string &s : v[i]) {
      a.push_back(s);
    }
  }

  for (int i = 0; i < K; i++) {
    pw[0][i] = 1;
  }

  for (int i = 1; i < M; i++) {
    for (int j = 0; j < K; j++) {
      pw[i][j] = (pw[i - 1][j] * primes[j]) % MOD;
    }
  }

  for (int i = 0; i < 26; i++) {
    t[i].max_load_factor(0.25);
    t[i].rehash(M);
  }

  ll ans = 0;
  for (string s : a) {
    debug(s);

    int p[26] = {0};
    for (char c : s) {
      p[c - 'a']++;
    }

    for (int j = 0; j < 26; j++) {
      if (p[j]) {
        ans += t[j][make_pair(0LL, 0LL)];
      }
    }

    ll h[K] = {0};
    for (int i = sz(s) - 1; i > 0; i--) {
      for (int j = 0; j < K; j++) {
        h[j] = (h[j] + (s[i] - 'a' + 1) * pw[sz(s) - i - 1][j]) % MOD;
      }

      p[s[i] - 'a']--;
      auto curH = make_pair(h[0], h[1]);
      for (int j = 0; j < 26; j++) {
        if (p[j]) {
          ans += t[j][curH];
        }
      }
    }

    t[s.front() - 'a'][make_pair(h[0], h[1])]++;
  }

  cout << ans << endl;

  return 0;
}