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

#include <bits/stdc++.h>
using namespace std;

typedef unsigned long long ull;
typedef long long ll;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;
typedef pair<double, double> pdd;
typedef vector<ll> vl;
typedef vector<vector<ll>> vvl;
// typedef vector<vector<ll>> Graph;

const ll mod = 1e9 + 7;
// const ll mod = 998244353;

#define REP(i, n) for (ll i = 0; i < (ll)n; i++)
#define dump(x) cerr << #x << " = " << (x) << endl;
#define spa << " " <<
#define fi first
#define se second

template <class T> bool chmax(T &a, const T &b) {
  if (a < b) {
    a = b;
    return 1;
  }
  return 0;
}
template <class T> bool chmin(T &a, const T &b) {
  if (a > b) {
    a = b;
    return 1;
  }
  return 0;
}

template <class S, class T>
ostream &operator<<(ostream &os, const pair<S, T> v) {
  os << "(" << v.first << ", " << v.second << ")";
  return os;
}
template <class T> ostream &operator<<(ostream &os, const vector<T> v) {
  for (int i = 0; i < (int)v.size(); i++) {
    if (i > 0) {
      os << " ";
    }
    os << v[i];
  }
  return os;
}
template <class T> ostream &operator<<(ostream &os, const vector<vector<T>> v) {
  for (int i = 0; i < (int)v.size(); i++) {
    if (i > 0) {
      os << endl;
    }
    os << v[i];
  }
  return os;
}

struct Trie {
  char a_init; // 最初の文字
  int D;       // 文字の種類数
  int INVALID = -1;
  std::vector<std::vector<int>> child;
  // 子のid．サイズは（頂点数，文字の種類数）で，child[v][c]は「頂点vから文字cで下降したときの頂点のid」
  // 行き先に頂点がない場合はINVALIDが格納される
  using T_NODE = int;
  vector<T_NODE> v_info;
  vector<int> par;
  vector<int> par_edge;
  // v_info には適当に頂点vの性質をメモすることができる
  // 例えば add_word の info に1を入れると出現回数のカウントができる
  Trie(char a_init = 'a', int D = 26)
      : a_init(a_init), D(D), child(1, std::vector<int>(D, INVALID)),
        v_info(1) {}
  int add_word(string str, T_NODE info) {
    int now = 0; // 頂点のid
    for (auto &c : str) {
      if (child[now][c - a_init] == INVALID) {
        int sz = child.size();
        child[now][c - a_init] = sz;
        child.emplace_back(std::vector<int>(D, INVALID));
        sz++;
        v_info.resize(sz);
        par.resize(sz);
        par_edge.resize(sz);
      }
      int ch = child[now][c - a_init];
      par[ch] = now;
      par_edge[ch] = c - a_init;
      now = ch;
    }
    v_info[now] += info;
    return now;
  }

  // 操作はこの上でdfsなどをすればOK
};

template <typename T> void debug(vector<vector<T>> &v, ll h, ll w) {
  for (ll i = 0; i < h; i++) {
    cerr << v[i][0];
    for (ll j = 1; j < w; j++)
      cerr spa v[i][j];
    cerr << endl;
  }
};
template <typename T> void debug(vector<T> &v, ll n) {
  if (n != 0)
    cerr << v[0];
  for (ll i = 1; i < n; i++)
    cerr spa v[i];
  cerr << endl;
};

int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);

  ll N;
  cin >> N;
  vector<string> S(N);
  REP(i, N) cin >> S[i];
  REP(i, N) reverse(S[i].begin(), S[i].end());

  Trie trie('a', 26);
  trie.add_word("", 0);
  vl idx(N);
  REP(i, N) { idx[i] = trie.add_word(S[i], 1); }

  ll res = 0;
  REP(i, N) {
    ll now = idx[i];
    set<ll> st;
    while (now > 0) {
      st.insert(trie.par_edge[now]);
      now = trie.par[now];
      for (auto s : st) {
        if (trie.child[now][s] != -1)
          res += trie.v_info[trie.child[now][s]];
      }
    }
  }

  cout << res - N << endl;

  return 0;
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;

typedef unsigned long long ull;
typedef long long ll;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;
typedef pair<double, double> pdd;
typedef vector<ll> vl;
typedef vector<vector<ll>> vvl;
// typedef vector<vector<ll>> Graph;

const ll mod = 1e9 + 7;
// const ll mod = 998244353;

#define REP(i, n) for (ll i = 0; i < (ll)n; i++)
#define dump(x) cerr << #x << " = " << (x) << endl;
#define spa << " " <<
#define fi first
#define se second

template <class T> bool chmax(T &a, const T &b) {
  if (a < b) {
    a = b;
    return 1;
  }
  return 0;
}
template <class T> bool chmin(T &a, const T &b) {
  if (a > b) {
    a = b;
    return 1;
  }
  return 0;
}

template <class S, class T>
ostream &operator<<(ostream &os, const pair<S, T> v) {
  os << "(" << v.first << ", " << v.second << ")";
  return os;
}
template <class T> ostream &operator<<(ostream &os, const vector<T> v) {
  for (int i = 0; i < (int)v.size(); i++) {
    if (i > 0) {
      os << " ";
    }
    os << v[i];
  }
  return os;
}
template <class T> ostream &operator<<(ostream &os, const vector<vector<T>> v) {
  for (int i = 0; i < (int)v.size(); i++) {
    if (i > 0) {
      os << endl;
    }
    os << v[i];
  }
  return os;
}

template <typename T> void debug(vector<vector<T>> &v, ll h, ll w) {
  for (ll i = 0; i < h; i++) {
    cerr << v[i][0];
    for (ll j = 1; j < w; j++)
      cerr spa v[i][j];
    cerr << endl;
  }
};
template <typename T> void debug(vector<T> &v, ll n) {
  if (n != 0)
    cerr << v[0];
  for (ll i = 1; i < n; i++)
    cerr spa v[i];
  cerr << endl;
};

struct rolling_hash {
  // !!! ハッシュはullで受け取ること !!!
  // 文字列の中に0があると，その値が影響しなくなって壊れるので注意！
  // generate_base
  // は結構遅いので，何回も呼び出す必要がある場合は何度も呼ばない方が良い
  typedef unsigned int ui;
  typedef unsigned long long ull;

  const ull MASK30 = (1ull << 30) - 1;
  const ull MASK31 = (1ull << 31) - 1;
  const ull MOD = (1ull << 61) - 1;
  const ull POSITIVISER = MOD * ((1ull << 3) - 1);
  ui base;
  int n;
  vector<ull> pw, hs;

  rolling_hash() {}

  rolling_hash(const string &s, int bs) {
    // set base
    base = (ui)bs;

    n = s.size();
    hs.assign(n + 1, 0);
    pw.assign(n + 1, 0);
    hs[0] = 0;
    pw[0] = 1;
    for (int i = 0; i < n; i++) {
      pw[i + 1] = calc_mod(mul(pw[i], base));
      hs[i + 1] = calc_mod(mul(hs[i], base) + s[i]);
    }
  }

  ull hash(int l, int r) {
    return calc_mod(hs[r] + POSITIVISER - mul(hs[l], pw[r - l]));
  }

  ull add_hash(ull r, int d) {
    ull hh = calc_mod(mul(r, base) + 'a' + d);
    return hh;
  }

  ull mul(ull l, ull r) {
    ull lu = l >> 31;
    ull ld = l & MASK31;
    ull ru = r >> 31;
    ull rd = r & MASK31;
    ull middle_bit = ld * ru + lu * rd;
    return ((lu * ru) << 1) + ld * rd + ((middle_bit & MASK30) << 31) +
           (middle_bit >> 30);
  }

  ull mul(ull l, ui r) {
    ull lu = l >> 31;
    ull rd = r & MASK31;
    ull middle_bit = lu * rd;
    return (l & MASK31) * rd + ((middle_bit & MASK30) << 31) +
           (middle_bit >> 30);
  }

  ull calc_mod(ull val) {
    val = (val & MOD) + (val >> 61);
    if (val > MOD)
      val -= MOD;
    return val;
  }
};

int generate_base() {
  random_device seed_gen;
  mt19937_64 rnd(seed_gen());
  uniform_int_distribution<int> dist(129, INT_MAX);
  int base = dist(rnd);
  return base;
}

int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);

  ll N;
  cin >> N;
  vector<string> S(N);
  REP(i, N) cin >> S[i];
  REP(i, N) reverse(S[i].begin(), S[i].end());

  int base = generate_base();
  vector<rolling_hash> R;
  REP(i, N) {
    rolling_hash rh(S[i], base);
    R.emplace_back(rh);
  }

  map<ull, ll> mp;
  REP(i, N) {
    ll sz = S[i].size();
    ull h = R[i].hash(0, sz);
    mp[h]++;
  }

  vector<ull> hs;
  ll res = 0;
  REP(i, N) {
    ll sz = S[i].size();
    vl cnt(26, 0);
    REP(j, sz) { cnt[S[i][j] - 'a']++; }
    REP(j, sz) {
      ull hh = R[i].hash(0, j);
      REP(k, 26) {
        if (cnt[k] > 0) {
          ull h = R[i].add_hash(hh, k);
          res += mp[h];
        }
      }
      cnt[S[i][j] - 'a']--;
    }
  }

  res -= N;
  cout << res << endl;

  return 0;
}