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

//@formatter:off
#include <bits/stdc++.h>
#define rep(i, n) for (int i = 0; i < int(n); ++i)
#define rrep(i, n) for (int i = int(n) - 1; i >= 0; i--)
#define rep2(i, s, n) for (int i = int(s); i < int(n); ++i)
#define all(a) a.begin(), a.end()
#define rall(a) a.rbegin(), a.rend()
#define pb push_back
#define eb emplace_back
#define vi vector<int>
#define vvi vector<vector<int>>
#define vl vector<ll>
#define vvl vector<vector<ll>>
#define vd vector<double>
#define vvd vector<vector<double>>
#define vs vector<string>
#define vc vector<char>
#define vvc vector<vector<char>>
#define vb vector<bool>
#define vvb vector<vector<bool>>
#define vp vector<P>
#define vvp vector<vector<P>>
using namespace std;
using ll = long long;
using P = pair<int, int>;
using LP = pair<ll, ll>;
template <class S, class T> istream &operator>>(istream &is, pair<S, T> &p) {
  return is >> p.first >> p.second;
}
template <class S, class T>
ostream &operator<<(ostream &os, const pair<S, T> &p) {
  return os << '{' << p.first << "," << p.second << '}';
}
template <class T> istream &operator>>(istream &is, vector<T> &v) {
  for (T &t : v) {
    is >> t;
  }
  return is;
}
template <class T> ostream &operator<<(ostream &os, const vector<T> &v) {
  os << '[';
  rep(i, v.size()) os << v[i] << (i == int(v.size() - 1) ? ']' : ',');
  return os;
}
void Yes(bool b) { cout << (b ? "Yes" : "No") << '\\n'; }
void YES(bool b) { cout << (b ? "YES" : "NO") << '\\n'; }
template <class T> bool chmin(T &a, T b) {
  if (a > b) {
    a = b;
    return true;
  }
  return false;
}
template <class T> bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}
const int inf = 1001001001;
const ll linf = 1001001001001001001;
//@formatter:on

using ull = unsigned long long;
const ull mod = (1ul << 61) - 1;
const ull mask30 = (1ul << 30) - 1;
const ull mask31 = (1ul << 31) - 1;
const ull mask61 = mod;

ull calc_mod(ull x) {
  ull xu = x >> 61;
  ull xd = x & mask61;
  ull res = xu + xd;
  if (res >= mod)
    res -= mod;
  return res;
}

// a*b mod 2^61-1
ull mul(ull a, ull b) {
  ull au = a >> 31;
  ull ad = a & mask31;
  ull bu = b >> 31;
  ull bd = b & mask31;
  ull mid = ad * bu + au * bd;
  ull midu = mid >> 30;
  ull midd = mid & mask30;
  return au * bu * 2 + midu + (midd << 31) + ad * bd;
}

ull base1;
ull base2;

class rolling_hash {
  int n;
  string s;
  vector<ull> hash1, hash2, pow1, pow2;

  void init() {
    hash1.assign(n + 1, 0);
    hash2.assign(n + 1, 0);
    pow1.assign(n + 1, 1);
    pow2.assign(n + 1, 1);
    rep(i, n) {
      hash1[i + 1] = calc_mod(mul(hash1[i], base1) + s[i]);
      hash2[i + 1] = calc_mod(mul(hash2[i], base2) + s[i]);
      pow1[i + 1] = calc_mod(mul(pow1[i], base1));
      pow2[i + 1] = calc_mod(mul(pow2[i], base2));
    }
  }

public:
  rolling_hash(string s) : s(s), n(s.size()) { init(); }
  // return hash of [l,r) of S
  pair<ull, ull> get(int l, int r) {
    ll res1 = calc_mod(hash1[r] + mod * 4 - mul(hash1[l], pow1[r - l]));
    ll res2 = calc_mod(hash2[r] + mod * 4 - mul(hash2[l], pow2[r - l]));
    return make_pair(res1, res2);
  }
  // return hash of T
  pair<ull, ull> get(string t) {
    ull ht1 = 0, ht2 = 0;
    rep(i, t.size()) {
      ht1 = calc_mod(mul(ht1, base1) + t[i]);
      ht2 = calc_mod(mul(ht2, base2) + t[i]);
    }
    return make_pair(ht1, ht2);
  }
  int count(string t) {
    if (t.size() > n)
      return 0;
    pair<ull, ull> ht = get(t);
    int res = 0;
    rep(i, n - t.size() + 1) {
      if (get(i, i + t.size()) == ht)
        res++;
    }
    return res;
  }
};

int main() {
  ios::sync_with_stdio(false);
  cin.tie(nullptr);
  cout.tie(nullptr);
  random_device rnd;
  mt19937_64 mt(rnd());
  base1 = calc_mod(mt());
  base2 = calc_mod(mt());
  while (base1 < 2 || base1 > mod - 2)
    base1 = calc_mod(mt());
  while (base2 < 2 || base2 > mod - 2)
    base2 = calc_mod(mt());
  int n;
  cin >> n;
  vs v(n);
  cin >> v;
  ll ans = 0;
  map<string, int> mp;
  rep(i, n) mp[v[i]]++;
  for (auto [f, s] : mp)
    ans += (ll)s * (s - 1) / 2;
  rep(i, n) reverse(all(v[i]));
  sort(all(v),
       [](const string &a, const string &b) { return a.size() < b.size(); });
  vector<map<pair<ull, ull>, int>> cnt(26);
  for (string &s : v) {
    vi ex(26);
    for (char &c : s)
      ex[c - 'a']++;
    rolling_hash rl(s);
    int sz = 0;
    while (sz < s.size() - 1) {
      auto p = rl.get(0, sz);
      rep(i, 26) if (ex[i] > 0) {
        if (cnt[i].count(p))
          ans += cnt[i][p];
      }
      ex[s[sz] - 'a']--;
      sz++;
    }
    cnt[s.back() - 'a'][rl.get(0, s.size() - 1)]++;
  }
  cout << ans << endl;
}


Error: TLE

Fixed submission:

//@formatter:off
#include <bits/stdc++.h>
#define rep(i, n) for (int i = 0; i < int(n); ++i)
#define rrep(i, n) for (int i = int(n) - 1; i >= 0; i--)
#define rep2(i, s, n) for (int i = int(s); i < int(n); ++i)
#define all(a) a.begin(), a.end()
#define rall(a) a.rbegin(), a.rend()
#define pb push_back
#define eb emplace_back
#define vi vector<int>
#define vvi vector<vector<int>>
#define vl vector<ll>
#define vvl vector<vector<ll>>
#define vd vector<double>
#define vvd vector<vector<double>>
#define vs vector<string>
#define vc vector<char>
#define vvc vector<vector<char>>
#define vb vector<bool>
#define vvb vector<vector<bool>>
#define vp vector<P>
#define vvp vector<vector<P>>
using namespace std;
using ll = long long;
using P = pair<int, int>;
using LP = pair<ll, ll>;
template <class S, class T> istream &operator>>(istream &is, pair<S, T> &p) {
  return is >> p.first >> p.second;
}
template <class S, class T>
ostream &operator<<(ostream &os, const pair<S, T> &p) {
  return os << '{' << p.first << "," << p.second << '}';
}
template <class T> istream &operator>>(istream &is, vector<T> &v) {
  for (T &t : v) {
    is >> t;
  }
  return is;
}
template <class T> ostream &operator<<(ostream &os, const vector<T> &v) {
  os << '[';
  rep(i, v.size()) os << v[i] << (i == int(v.size() - 1) ? ']' : ',');
  return os;
}
void Yes(bool b) { cout << (b ? "Yes" : "No") << '\\n'; }
void YES(bool b) { cout << (b ? "YES" : "NO") << '\\n'; }
template <class T> bool chmin(T &a, T b) {
  if (a > b) {
    a = b;
    return true;
  }
  return false;
}
template <class T> bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}
const int inf = 1001001001;
const ll linf = 1001001001001001001;
//@formatter:on

int main() {
  ios::sync_with_stdio(false);
  cin.tie(nullptr);
  cout.tie(nullptr);
  int n;
  cin >> n;
  vs v(n);
  cin >> v;
  ll ans = 0;
  map<string, int> mp;
  rep(i, n) mp[v[i]]++;
  for (auto [f, s] : mp)
    ans += (ll)s * (s - 1) / 2;
  rep(i, n) reverse(all(v[i]));
  sort(all(v),
       [](const string &a, const string &b) { return a.size() < b.size(); });
  vector<map<string, int>> cnt(26);
  for (string &s : v) {
    vi ex(26);
    for (char &c : s)
      ex[c - 'a']++;
    string t = "";
    while (t.size() < s.size() - 1) {
      rep(i, 26) if (ex[i] > 0) {
        if (cnt[i].count(t))
          ans += cnt[i][t];
      }
      char c = s[t.size()];
      ex[c - 'a']--;
      t.pb(c);
    }
    string now = s;
    now.pop_back();
    char c = s.back();
    cnt[c - 'a'][now]++;
  }
  cout << ans << endl;
}