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

#include <bits/stdc++.h>
#define endl "\\n"
using namespace std;
typedef long long ll;
typedef pair<ll, ll> l_l;
typedef pair<int, int> i_i;
template <class T> inline bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}

template <class T> inline bool chmin(T &a, T b) {
  if (a > b) {
    a = b;
    return true;
  }
  return false;
}

const long long INF = 1e18;
// const ll mod = 1000000007;
ll num[1050000][26];
template <int char_size> struct TrieNode {
  int nxt[char_size];

  int exist;
  vector<int> accept;

  TrieNode() : exist(0) { memset(nxt, -1, sizeof(nxt)); }
};

template <int char_size, int margin> struct Trie {
  using Node = TrieNode<char_size>;

  vector<Node> nodes;
  int root;

  Trie() : root(0) { nodes.push_back(Node()); }

  void update_direct(int node, int id) { nodes[node].accept.push_back(id); }

  void update_child(int node, int child, int id) { ++nodes[node].exist; }

  void add(const string &str, int str_index, int node_index, int id) {
    if (str_index + 1 == str.size()) {
      ll c = str[str_index] - margin;
      num[node_index][c]++;
    }
    if (str_index == str.size()) {
      update_direct(node_index, id);
    } else {
      const int c = str[str_index] - margin;
      if (nodes[node_index].nxt[c] == -1) {
        nodes[node_index].nxt[c] = (int)nodes.size();
        nodes.push_back(Node());
      }
      add(str, str_index + 1, nodes[node_index].nxt[c], id);
      update_child(node_index, nodes[node_index].nxt[c], id);
    }
  }

  void add(const string &str, int id) { add(str, 0, 0, id); }

  void add(const string &str) { add(str, nodes[0].exist); }

  void query(const string &str, const function<void(int)> &f, int str_index,
             int node_index) {
    for (auto &idx : nodes[node_index].accept)
      f(idx);
    if (str_index == str.size()) {
      return;
    } else {
      const int c = str[str_index] - margin;
      if (nodes[node_index].nxt[c] == -1)
        return;
      query(str, f, str_index + 1, nodes[node_index].nxt[c]);
    }
  }

  void query(const string &str, const function<void(int)> &f) {
    query(str, f, 0, 0);
  }

  int count() const { return (nodes[0].exist); }

  int size() const { return ((int)nodes.size()); }
};
ll ans;
ll N;
vector<string> S;
ll tmpnum[26];
int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);
  Trie<26, 'a'> trie;
  cin >> N;
  S.resize(N);
  for (int i = 0; i < N; i++) {
    cin >> S[i];
    reverse(S[i].begin(), S[i].end());
    trie.add(S[i]);
  }
  /*
  for(int i = 0; i < trie.size(); i++) {
      for(int j = 0; j < 26; j++) {
          if(num[i][j]) cerr << i << " " << j << " " << num[i][j] << endl;
      }
  }
  */
  for (int i = 0; i < N; i++) {
    ll node = 0;
    for (ll j = 0; j < S[i].size(); j++) {
      tmpnum[(int)(S[i][j] - 'a')]++;
    }
    for (int j = 0; j < S[i].size(); j++) {
      for (int k = 0; k < 26; k++) {
        if (tmpnum[k]) {
          // cerr << i << " " << j << " " << k << " " << tmpnum[k] << endl;
          ans += num[node][k];
        }
      }
      node = trie.nodes[node].nxt[(int)(S[i][j] - 'a')];
      tmpnum[(int)(S[i][j] - 'a')]--;
    }
  }
  cout << ans - N << endl;
  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
#define endl "\\n"
using namespace std;
typedef long long ll;
typedef pair<ll, ll> l_l;
typedef pair<int, int> i_i;
template <class T> inline bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}

template <class T> inline bool chmin(T &a, T b) {
  if (a > b) {
    a = b;
    return true;
  }
  return false;
}

const long long INF = 1e18;
// const ll mod = 1000000007;
ll ans;
ll N;
vector<string> S;
template <class Z> Z rng(Z a, Z b) {
  static mt19937 mt(chrono::steady_clock::now().time_since_epoch().count());
  return uniform_int_distribution<Z>(a, b - 1)(mt);
}

struct RollingHash {
  static constexpr uint64_t P0 = 4e9 + 7;
  static constexpr uint64_t P1 = 4e9 + 9;
  static uint64_t B0;
  static uint64_t B1;
  static vector<uint64_t> powB0;
  static vector<uint64_t> powB1;
  const int n;
  vector<uint64_t> h0;
  vector<uint64_t> h1;
  template <class Itr>
  RollingHash(Itr first, Itr last)
      : n(distance(first, last)), h0(n + 1), h1(n + 1) {
    for (int i = 0; i < n; ++i, ++first) {
      h0[i + 1] = (h0[i] * B0 + *first) % P0;
      h1[i + 1] = (h1[i] * B1 + *first) % P1;
    }
    while (powB0.size() <= n) {
      powB0.push_back(powB0.back() * B0 % P0);
      powB1.push_back(powB1.back() * B1 % P1);
    }
  }
  uint64_t get0(int l, int r) {
    return (h0[r] + (P0 - h0[l]) * powB0[r - l]) % P0;
  }
  uint64_t get1(int l, int r) {
    return (h1[r] + (P1 - h1[l]) * powB1[r - l]) % P1;
  }
  pair<uint64_t, uint64_t> get(int l, int r) {
    return {get0(l, r), get1(l, r)};
  }
  bool eq(int l0, int r0, int l1, int r1) {
    return get0(l0, r0) == get0(l1, r1) and get1(l0, r0) == get1(l1, r1);
  }
  template <class Itr> static uint64_t get0(Itr first, Itr last) {
    uint64_t res = 0;
    while (first != last) {
      res = (res * B0 + *first++) % P0;
    }
    return res;
  }
  template <class Itr> static uint64_t get1(Itr first, Itr last) {
    uint64_t res = 0;
    while (first != last) {
      res = (res * B1 + *first++) % P1;
    }
    return res;
  }
};
uint64_t RollingHash::B0 = rng<unsigned>(1, RollingHash::P0);
uint64_t RollingHash::B1 = rng<unsigned>(1, RollingHash::P1);
vector<uint64_t> RollingHash::powB0{1};
vector<uint64_t> RollingHash::powB1{1};
typedef pair<uint64_t, uint64_t> uu;
int main() {
  map<pair<uint64_t, uint64_t>, ll> mp[26];
  cin >> N;
  S.resize(N);
  for (int i = 0; i < N; i++)
    cin >> S[i];
  sort(S.begin(), S.end(),
       [](string a, string b) { return a.size() < b.size(); });
  for (int i = 0; i < N; i++) {
    RollingHash RHS(S[i].begin(), S[i].end());
    int bits = 0;
    for (int j = 0; j < S[i].size(); j++) {
      uu Hash = RHS.get(j + 1, S[i].size());
      bits |= (1 << ((int)(S[i][j] - 'a')));
      for (int k = 0; k < 26; k++) {
        if (bits & (1 << k))
          ans += mp[k][Hash];
      }
    }
    uu Hash = RHS.get(1, S[i].size());
    ll idx = (int)(S[i][0] - 'a');
    mp[idx][Hash]++;
  }
  cout << ans << endl;
  return 0;
}