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

#include <bits/stdc++.h>
using namespace std;
#define ll long long
#define mp make_pair
#define pb push_back
typedef pair<int, int> P;

const int K = 26; // size of used alphabets.
const int N = 2e5 + 5;

struct Vertex {
  int nxt[K];
  bool isLeaf = false;
  Vertex() { fill(nxt, nxt + K, -1); }
};

vector<Vertex> trie(1);
vector<int> path[N];
string s[N];
int res;

void add(int id, const string &s) {
  int v = 0;
  for (const char &ch : s) {
    int c = ch - 'a';
    path[id].pb(v);
    if (trie[v].nxt[c] == -1) {
      trie[v].nxt[c] = trie.size();
      trie.emplace_back();
    }
    v = trie[v].nxt[c];
  }
  trie[v].isLeaf = true;
}

int main() {
  ios_base::sync_with_stdio(0);
  cin.tie(0);
  int n;
  cin >> n;
  for (int i = 0; i < n; ++i) {
    cin >> s[i];
    reverse(s[i].begin(), s[i].end());
    add(i, s[i]);
  }
  for (int i = 0; i < n; ++i) {
    set<char> c;
    for (int j = int(s[i].size()) - 1; ~j; --j) {
      c.insert(s[i][j]);
      int v = path[i][j];
      for (char ch : c)
        res +=
            trie[v].nxt[ch - 'a'] != -1 && trie[trie[v].nxt[ch - 'a']].isLeaf;
    }
  }
  cout << res - n << '\\n';
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
#define ll long long
#define mp make_pair
#define pb push_back
typedef pair<int, int> P;

const int N = 1e6 + 5;
const ll MOD = 1e9 + 9;
const ll MOD2 = 998244353;
mt19937_64 rng(chrono::steady_clock::now().time_since_epoch().count());
ll base = 311;
ll base2 = rng() % 300;
ll pb[N], pb2[N];
struct double_hash {
  int sz;
  vector<ll> pref[2], suff[2];
  double_hash(string s) {
    sz = s.size();
    for (int i = 0; i < 2; ++i)
      pref[i].resize(sz + 5, 0), suff[i].resize(sz + 5, 0);
    pref[0][0] = pref[1][0] = s[0] - 'a' + 1LL;
    for (int i = 1; i < sz; ++i) {
      pref[0][i] = (pref[0][i - 1] * base % MOD + (s[i] - 'a' + 1LL)) % MOD;
      pref[1][i] = (pref[1][i - 1] * base2 % MOD2 + (s[i] - 'a' + 1LL)) % MOD2;
    }
    reverse(s.begin(), s.end());
    suff[0][0] = suff[1][0] = s[0] - 'a' + 1LL;
    for (int i = 1; i < sz; ++i) {
      suff[0][i] =
          (suff[0][i - 1] * base % MOD + (s[i] - 'a' + 1LL) % MOD) % MOD;
      suff[1][i] =
          (suff[1][i - 1] * base2 % MOD2 + (s[i] - 'a' + 1LL) % MOD2) % MOD2;
    }
  }
  pair<ll, ll> get_prefix(int l, int r) {
    ll val =
        (pref[0][r] - (l ? pref[0][l - 1] : 0LL) * pb[r - l + 1] + MOD * MOD) %
        MOD;
    ll val2 = (pref[1][r] - (l ? pref[1][l - 1] : 0LL) * pb2[r - l + 1] +
               MOD2 * MOD2) %
              MOD2;
    return pair<ll, ll>(val, val2);
  }
  pair<ll, ll> get_suffix(int l, int r) {
    l = sz - l - 1, r = sz - r - 1;
    swap(l, r);
    ll val =
        (suff[0][r] - (l ? suff[0][l - 1] : 0LL) * pb[r - l + 1] + MOD * MOD) %
        MOD;
    ll val2 = (suff[1][r] - (l ? suff[1][l - 1] : 0LL) * pb2[r - l + 1] +
               MOD2 * MOD2) %
              MOD2;
    return pair<ll, ll>(val, val2);
  }
};

int n;
vector<string> a;
map<pair<char, pair<ll, ll>>, int> Map;

int main() {
  ios_base::sync_with_stdio(0);
  cin.tie(0);
  pb[0] = pb2[0] = 1LL;
  for (int i = 1; i < N; ++i) {
    pb[i] = pb[i - 1] * base % MOD;
    pb2[i] = pb2[i - 1] * base2 % MOD2;
  }
  ll res = 0;
  cin >> n;
  for (int i = 0; i < n; ++i) {
    string s;
    cin >> s;
    a.push_back(s);
    int len = s.size();
    set<char> c;
    double_hash hash_s = double_hash(s);
    if (len == 1)
      res += Map[mp(s[0], mp(-1, -1))];
    else
      res += Map[mp(s[0], hash_s.get_prefix(1, len - 1))];
    for (int i = 0; i < len; ++i) {
      c.insert(s[i]);
      for (char x : c)
        if (i == len - 1)
          Map[mp(x, mp(-1, -1))]++;
        else
          Map[mp(x, hash_s.get_prefix(i + 1, len - 1))]++;
    }
  }
  reverse(a.begin(), a.end());
  Map.clear();
  for (int i = 0; i < n; ++i) {
    string s = a[i];
    int len = s.size();
    set<char> c;
    double_hash hash_s = double_hash(s);
    if (len == 1)
      res += Map[mp(s[0], mp(-1, -1))];
    else
      res += Map[mp(s[0], hash_s.get_prefix(1, len - 1))];
    for (int i = 0; i < len; ++i) {
      c.insert(s[i]);
      for (char x : c)
        if (i == len - 1)
          Map[mp(x, mp(-1, -1))]++;
        else
          Map[mp(x, hash_s.get_prefix(i + 1, len - 1))]++;
    }
  }
  cout << res << '\\n';
}