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

#include <bits/stdc++.h>

using namespace std;

int n;

int main() {
  ios::sync_with_stdio(false);
  cin.tie(0);
  cout.tie(0);
  cin >> n;
  vector<string> s(n);
  for (int i = 0; i < n; i++) {
    cin >> s[i];
    reverse(s[i].begin(), s[i].end());
  }
  sort(s.begin(), s.end(),
       [&](string a, string b) { return a.length() > b.length(); });
  vector<int> leaf(1);
  vector<vector<int>> to(1, vector<int>(26, -1));
  vector<vector<int>> sum(1, vector<int>(26));
  long long res = 0;
  for (string p : s) {
    int sz = p.length();
    int cur = 0;
    vector<int> path;
    path.push_back(cur);
    for (int i = 0; i < sz; i++) {
      int c = p[i] - 'a';
      if (to[cur][c] == -1) {
        to[cur][c] = leaf.size();
        leaf.push_back(0);
        to.push_back(vector<int>(26, -1));
        sum.push_back(vector<int>(26));
        cur = leaf.size() - 1;
      } else {
        cur = to[cur][c];
      }
      path.push_back(cur);
    }
    int root = path[sz - 1];
    int last_char = p.back() - 'a';
    res += sum[root][last_char];
    for (int i = 1; i <= sz; i++) {
      int node = path[i];
      int par = path[i - 1];
      vector<int> rev = sum[node];
      for (int j = 0; j < 26; j++)
        if (to[par][j] == node)
          rev[j] = leaf[node];
      for (int j = 0; j < 26; j++)
        sum[par][j] -= rev[j];
    }
    for (int i = 0; i <= sz; i++) {
      int node = path[i];
      leaf[node]++;
    }
    for (int i = sz; i > 0; i--) {
      int par = path[i - 1];
      int node = path[i];
      vector<int> add = sum[node];
      for (int j = 0; j < 26; j++)
        if (to[par][j] == node)
          add[j] = leaf[node];
      for (int j = 0; j < 26; j++)
        sum[par][j] += add[j];
    }
  }
  cout << res;
  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>

using namespace std;

inline int inv(int a, int md) {
  a %= md;
  if (a < 0)
    a += md;
  int b = md, u = 0, v = 1;
  while (a) {
    int t = b / a;
    b -= t * a;
    swap(a, b);
    u -= t * v;
    swap(u, v);
  }
  assert(b == 1);
  if (u < 0)
    u += md;
  return u;
}

class HASH {
public:
  ;
  int a = 1e9 + 7;
  int b = 998244353;
  int x = 0, y = 0;
  HASH(){};
  HASH(int _x, int _y) : x(_x), y(_y) {
    if (x < 0)
      x += a;
    if (x >= a)
      x -= a;
    if (y < 0)
      y += b;
    if (y >= b)
      y -= b;
  };

  HASH operator+(const HASH &c) const { return HASH(x + c.x, y + c.y); }
  HASH operator-(const HASH &c) const { return HASH(x - c.x, y - c.y); }
  HASH operator*(const HASH &c) const {
    return HASH(((long long)x * c.x) % a, ((long long)y * c.y) % b);
  }
  HASH operator*(int c) const {
    return HASH(((long long)x * c) % a, ((long long)y * c) % b);
  }
  HASH operator/(const HASH &c) const {
    HASH _c = HASH(inv(c.x, a), inv(c.y, b));
    return *this * _c;
  }
  HASH operator/(int c) const { return *this / HASH(c, c); }
  bool operator<(const HASH &c) const {
    return make_pair(x, y) < make_pair(c.x, c.y);
  }
};

class suffix_hash {
public:
  ;
  string s;
  vector<HASH> pref;
  HASH c;
  int n;
  vector<HASH> p;
  suffix_hash(string _s) : s(_s) {
    HASH c(97, 31);
    n = s.length();
    pref.resize(n + 1);
    p.resize(n + 1);
    p[0] = HASH(1, 1);
    for (int i = 1; i < n + 1; i++) {
      p[i] = p[i - 1] * c;
    }
    for (int i = 1; i < n + 1; i++) {
      pref[i] = pref[i - 1] + p[i] * (s[i - 1] - 'a' + 3);
    }
  }
  long long get(int l, int r) {
    int md = 1e9 + 7;
    HASH cc = (pref[r] - pref[l]) / p[l];
    return 1LL * cc.x * md + cc.y;
  }
};

int n;

int main() {
  ios::sync_with_stdio(false);
  cin.tie(0);
  cout.tie(0);
  cin >> n;
  vector<string> s(n);
  for (int i = 0; i < n; i++) {
    cin >> s[i];
    reverse(s[i].begin(), s[i].end());
  }
  sort(s.begin(), s.end(),
       [&](string a, string b) { return a.length() > b.length(); });
  map<pair<long long, char>, int> cnt;
  long long res = 0;
  for (int i = 0; i < n; i++) {
    vector<int> last(26, -1);
    for (int j = 0; j < (int)s[i].length(); j++) {
      last[s[i][j] - 'a'] = j;
    }
    suffix_hash cur(s[i]);
    long long c = cur.get(0, s[i].length() - 1);
    res += cnt[make_pair(c, s[i].back())];
    for (int j = 0; j < (int)s[i].length(); j++) {
      for (int u = 0; u < 26; u++) {
        char t = 'a' + u;
        long long p = cur.get(0, j + 1);
        if (last[u] <= j)
          continue;
        cnt[make_pair(p, t)]++;
      }
    }
    for (int j = 0; j < 26; j++)
      if (last[j] > -1)
        cnt[make_pair(0, 'a' + j)]++;
  }
  cout << res;
  return 0;
}