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

#include <algorithm>
#include <iostream>
#include <string>
#include <vector>

using namespace std;

using ll = long long;
#define REP(i, n) for (int(i) = 0; (i) < (n); (i)++)
#define all(x) (x).begin(), (x).end()
struct Trie {
  struct Node {
    int c;
    bool accept;
    vector<int> nextid;
    Node(int c) : c(c), accept(false) { nextid.assign(26, -1); }
  };
  vector<Node> g;
  Trie() { g.emplace_back(0); }
  bool match(int tmp, int c) {
    //    assert(tmp < g.size());
    return (g[tmp].nextid[c] != -1) && g[g[tmp].nextid[c]].accept;
  }
  void insert(string &s) {
    int tmp = 0;
    for (int i = 0; i < s.size(); i++) {
      int &nid = g[tmp].nextid[s[i] - 'a'];
      if (nid == -1) {
        nid = g.size();
        g.emplace_back(s[i] - 'a');
      }
      tmp = nid;
    }
    g[tmp].accept = true;
  }
};

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

  int n;
  cin >> n;
  vector<string> a(n);
  for (int i = 0; i < n; i++) {
    cin >> a[i];
    reverse(all(a[i]));
  }
  sort(a.begin(), a.end(),
       [](string &s, string &t) { return s.size() < t.size(); });

  Trie trie;
  ll ans = 0;
  for (int i = 0; i < n; i++) {
    string s = a[i];
    int sz = s.size();
    vector<int> last(26, -1);
    for (int j = sz - 1; j >= 0; j--) {
      if (last[s[j] - 'a'] == -1) {
        last[s[j] - 'a'] = j;
      }
    }
    int tmp = 0;
    for (int j = 0; j < sz; j++) {
      for (int c = 0; c < 26; c++) {
        if (j <= last[c] && trie.match(tmp, c)) {
          ans++;
        }
      }
      int nid = trie.g[tmp].nextid[s[j] - 'a'];
      if (nid == -1)
        break;
      else
        tmp = nid;
    }
    trie.insert(s);
  }
  cout << ans << '\\n';
  return 0;
}


Error: TLE

Fixed submission:

#pragma GCC target("avx2")
#pragma GCC optimize("O3")
#pragma GCC optimize("unroll-loops")
#include <bits/stdc++.h>
using namespace std;

using ll = long long;
#define REP(i, n) for (int i = 0; i < n; i++)

#define maxn 200020
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 == 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()); }
};

int n;
string s[maxn];

bool comp(string &s, string &t) {
  if (s.size() != t.size()) {
    return s.size() > t.size();
  } else
    return s < t;
}
int main(void) {
  cin.tie(0);
  ios::sync_with_stdio(false);

  cin >> n;
  REP(i, n) cin >> s[i];
  REP(i, n) reverse(s[i].begin(), s[i].end());
  sort(s, s + n, comp);

  // REP(i,n){
  // 	cout << s[i] << '\\n';
  // }
  // return 0;
  ll ans = 0;
  // map<string,ll> M;
  set<string> S[26];
  // Trie<26,'a'> trie;
  for (int i = n - 1; i >= 0; i--) {
    int sz = s[i].size();
    string tmp = s[i];
    set<char> st;
    int id = s[i][0] - 'a';
    for (int j = sz - 1; j >= 0; j--) {
      st.insert(s[i][j]);
      for (char x : st) {
        tmp.back() = x;
        // ans += M[tmp];
        ans += (ll)(S[tmp[0] - 'a'].find(tmp) != S[tmp[0] - 'a'].end());
        // ans += trie.
      }
      tmp.erase(j);
      // S.insert(tmp);
    }
    // M[s[i]]++;
    S[s[i][0] - 'a'].insert(s[i]);
  }
  cout << ans << '\\n';
}