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

#pragma GCC target("avx2")
#pragma GCC optimize("unroll-loops")
#pragma GCC optimize("O3")
#include "bits/stdc++.h"
#include <random>
#include <unordered_map>
#include <unordered_set>
using namespace std;
typedef long long ll;
const ll MOD = /*1'000'000'007LL;*/ 998'244'353LL;
#define pb push_back
#define mp make_pair
#define all(x) (x).begin(), (x).end()
#define rep(i, n) for (int(i) = 0; (i) < (n); (i)++)
const int dx[4] = {1, 0, -1, 0};
const int dy[4] = {0, 1, 0, -1};

void operator+=(vector<int> &a, vector<int> &b) {
  rep(i, a.size()) a[i] += b[i];
}

template <int sigma, int a, typename T> struct Trie {
private:
  T e;
  struct Node {
    int chi[sigma];
    T x;
    Node(T e) {
      x = e;
      memset(chi, -1, sizeof(chi));
    }
  };
  vector<Node> nodes;

public:
  Trie(T _e) {
    e = _e;
    nodes.push_back(Node(e));
  }
  void add_node(int n, int c) {
    nodes[n].chi[c] = nodes.size();
    nodes.push_back(Node(e));
  }
  void add(string s, T y) {
    int n = 0;
    nodes[n].x += y;
    for (int i = 0; i < s.size(); i++) {
      if (nodes[n].chi[s[i] - a] == -1)
        add_node(n, s[i] - a);
      n = nodes[n].chi[s[i] - a];
      nodes[n].x += y;
    }
  }
  T count_prefix(string s) {
    int n = 0;
    for (int i = 0; i < s.size(); i++) {
      if (nodes[n].chi[s[i] - a] == -1)
        return e;
      n = nodes[n].chi[s[i] - a];
    }
    return nodes[n].x;
  }
};

int N;
vector<string> S[1000001];

Trie<26, 'a', vector<int>> trie(vector<int>(26, 0));
vector<pair<int, pair<int, int>>> query[1000001];

signed main() {
  cin >> N;
  rep(i, N) {
    string s;
    cin >> s;
    reverse(all(s));
    S[s.size()].pb(s);
  }

  ll ans = 0;

  vector<int> v(26, 0);
  for (int i = 1000000; i >= 1; i--) {
    rep(j, query[i - 1].size()) {
      v[query[i - 1][j].first] = 1;
      trie.add(S[query[i - 1][j].second.first][query[i - 1][j].second.second],
               v);
      v[query[i - 1][j].first] = 0;
      query[i - 1][j].first = 0;
    }
    rep(j, S[i].size()) {
      ans += trie.count_prefix(
          S[i][j].substr(0, S[i][j].size() - 1))[S[i][j].back() - 'a'];
      int p[26];
      memset(p, -1, sizeof(p));
      rep(l, S[i][j].size()) { p[S[i][j][l] - 'a'] = l; }
      rep(l, 26) {
        if (p[l] == i - 1) {
          v[l] = 1;
          trie.add(S[i][j], v);
          v[l] = 0;
        } else if (p[l] != -1)
          query[p[l]].pb({l, {i, j}});
      }
    }
  }

  cout << ans << endl;
}

Error: TLE

Fixed submission:

#pragma GCC target("avx2")
#pragma GCC optimize("unroll-loops")
#pragma GCC optimize("O3")
#include "bits/stdc++.h"
#include <random>
#include <unordered_map>
#include <unordered_set>
using namespace std;
typedef long long ll;
const ll MOD = /*1'000'000'007LL;*/ 998'244'353LL;
#define pb push_back
#define mp make_pair
#define all(x) (x).begin(), (x).end()
#define rep(i, n) for (int(i) = 0; (i) < (n); (i)++)
const int dx[4] = {1, 0, -1, 0};
const int dy[4] = {0, 1, 0, -1};

template <int sigma, int a> struct Trie {
private:
  struct Node {
    int chi[sigma];
    int siz = 0;
    Node() { memset(chi, -1, sizeof(chi)); }
  };
  vector<Node> nodes;

public:
  Trie() { nodes.push_back(Node()); }
  void add_node(int n, int c) {
    nodes[n].chi[c] = nodes.size();
    nodes.push_back(Node());
  }
  void add(string &s) {
    int n = 0;
    nodes[n].siz++;
    for (int i = 0; i < s.size(); i++) {
      if (nodes[n].chi[s[i] - a] == -1)
        add_node(n, s[i] - a);
      n = nodes[n].chi[s[i] - a];
      nodes[n].siz++;
    }
  }
  int count_prefix(string &s) {
    int n = 0;
    for (int i = 0; i < s.size(); i++) {
      if (nodes[n].chi[s[i] - a] == -1)
        return 0;
      n = nodes[n].chi[s[i] - a];
    }
    return nodes[n].siz;
  }
};

int N;
vector<string> S[1000001];

Trie<26, 'a'> trie[26];
vector<pair<int, pair<int, int>>> query[1000001];

signed main() {
  cin >> N;
  rep(i, N) {
    string s;
    cin >> s;
    reverse(all(s));
    S[s.size()].pb(s);
  }

  ll ans = 0;

  for (int i = 1000000; i >= 1; i--) {
    rep(j, query[i - 1].size()) {
      trie[query[i - 1][j].first].add(
          S[query[i - 1][j].second.first][query[i - 1][j].second.second]);
    }
    rep(j, S[i].size()) {
      string tmp = S[i][j].substr(0, S[i][j].size() - 1);
      ans += trie[S[i][j].back() - 'a'].count_prefix(tmp);
      int p[26];
      memset(p, -1, sizeof(p));
      rep(l, S[i][j].size()) { p[S[i][j][l] - 'a'] = l; }
      rep(l, 26) {
        if (p[l] == i - 1)
          trie[l].add(S[i][j]);
        else if (p[l] != -1)
          query[p[l]].pb({l, {i, j}});
      }
    }
  }

  cout << ans << endl;
}