Problem p02589 (C++) — buggy submission, verdict: Runtime Error

#include <bits/stdc++.h>
#define ll long long int
#define ld long double
#define f first
#define s second
#define pb push_back
#define eb emplace_back
#define mk make_pair
#define mt make_tuple
#define MOD 1000000007
#define fo(i, a, b) for (i = a; i < b; i++)
#define foe(i, a, b) for (i = a; i <= b; i++)
#define all(x) x.begin(), x.end()
#define vi vector<int>
#define vl vector<long long int>
#define pii pair<int, int>
#define pll pair<long long int, long long int>
#define vpii vector<pair<int, int>>
#define vpll vector<pair<long long int, long long int>>
#define boost                                                                  \\
  ios::sync_with_stdio(false);                                                 \\
  cin.tie(0)
using namespace std;
const int inf = 1e9 + 5;
const ll inf64 = 1e18 + 5;

const int MAX = 26;
int cnt[MAX];

struct Trie {
  Trie *po[26];
  vector<ll> v;
  ll cnt = 0;
};
Trie *newNode() {
  Trie *node = new Trie();
  node->cnt = 0;
  node->v.resize(26);
  for (int i = 0; i < 26; i++) {
    node->po[i] = NULL;
    node->v[i] = 0;
  }
  return node;
}
void insert(Trie *root, string &s, int ind, vector<vector<int>> &w) {
  if (ind == s.size())
    return;
  int x = s[ind] - 'a';

  if (root->po[x] == NULL)
    root->po[x] = newNode();
  root->po[x]->cnt++;
  for (int i = 0; i < 26; i++)
    root->po[x]->v[i] += (w[ind][i] > 0);

  insert(root->po[x], s, ind + 1, w);
}

ll qry1(Trie *root, string &s, int ind) {
  if (ind == s.size())
    return root->cnt;

  int x = s[ind] - 'a';
  if (root->po[x] == NULL)
    return 0;
  return qry1(root->po[x], s, ind + 1);
}

ll qry2(Trie *root, string &s, int ind, int req) {
  if (ind == s.size() - 1)
    return root->v[req];

  int x = s[ind] - 'a';
  if (root->po[x] == NULL)
    return 0;
  return qry2(root->po[x], s, ind + 1, req);
}

int main() {
  boost;
  int n, i;
  cin >> n;
  vector<string> v(n);

  string s;
  fo(i, 0, n) {
    cin >> s;
    reverse(all(s));
    v[i] = s;
  }

  sort(all(v), [&](string &s1, string &s2) {
    if (s1.size() == s2.size())
      return s1 < s2;
    return s1.size() > s2.size();
  });
  Trie *root = newNode();

  ll ans = 0;
  fo(i, 0, n) {
    if (v[i].size() > 1) {
      ll res = qry1(root, v[i], 0);
      ans += res;
      ans += qry2(root, v[i], 0, v[i].back() - 'a') - res;
    } else
      ans += cnt[v[i][0] - 'a'];

    s = v[i];
    int x = s.size();
    vector<vector<int>> w(x, vector<int>(26, 0));
    for (char c : s)
      w[0][c - 'a']++;

    bool done[26];
    memset(done, 0, sizeof done);

    for (char c : s) {
      if (!done[c - 'a']) {
        done[c - 'a'] = 1;
        cnt[c - 'a']++;
      }
    }

    for (int j = 0; j < s.size(); j++) {
      if (j)
        w[j] = w[j - 1];
      w[j][s[j] - 'a']--;
    }

    insert(root, s, 0, w);
  }
  cout << ans;
}


Error: WA

Fixed submission:

#include <bits/stdc++.h>
#define ll long long int
#define ld long double
#define f first
#define s second
#define pb push_back
#define eb emplace_back
#define mk make_pair
#define mt make_tuple
#define MOD 1000000007
#define fo(i, a, b) for (i = a; i < b; i++)
#define foe(i, a, b) for (i = a; i <= b; i++)
#define all(x) x.begin(), x.end()
#define vi vector<int>
#define vl vector<long long int>
#define pii pair<int, int>
#define pll pair<long long int, long long int>
#define vpii vector<pair<int, int>>
#define vpll vector<pair<long long int, long long int>>
#define boost                                                                  \\
  ios::sync_with_stdio(false);                                                 \\
  cin.tie(0)
using namespace std;
const int inf = 1e9 + 5;
const ll inf64 = 1e18 + 5;

const int MAX = 26;
int cnt[MAX];

struct Trie {
  Trie *po[26];
  vector<ll> v;
  ll cnt = 0;
};
Trie *newNode() {
  Trie *node = new Trie();
  node->cnt = 0;
  node->v.resize(26);
  for (int i = 0; i < 26; i++) {
    node->po[i] = NULL;
    node->v[i] = 0;
  }
  return node;
}
void insert(Trie *root, string &s, int ind, vector<vector<int>> &w) {
  if (ind == s.size())
    return;
  int x = s[ind] - 'a';

  if (root->po[x] == NULL)
    root->po[x] = newNode();
  root->po[x]->cnt++;
  for (int i = 0; i < 26; i++)
    root->po[x]->v[i] += (w[ind][i] > 0);

  insert(root->po[x], s, ind + 1, w);
}

ll qry1(Trie *root, string &s, int ind) {
  if (ind == s.size())
    return root->cnt;

  int x = s[ind] - 'a';
  if (root->po[x] == NULL)
    return 0;
  qry1(root->po[x], s, ind + 1);
}

ll qry2(Trie *root, string &s, int ind, int req) {
  if (ind == s.size() - 1)
    return root->v[req];

  int x = s[ind] - 'a';
  if (root->po[x] == NULL)
    return 0;
  qry2(root->po[x], s, ind + 1, req);
}

int main() {
  boost;
  int n, i;
  cin >> n;
  string s;
  vector<string> v(n);

  Trie *root = newNode();
  fo(i, 0, n) {
    cin >> s;
    reverse(all(s));

    v[i] = s;
    int x = s.size();
    vector<vector<int>> w(x, vector<int>(26, 0));
    for (char c : s)
      w[0][c - 'a']++;

    bool done[26];
    memset(done, 0, sizeof done);

    for (char c : s) {
      if (!done[c - 'a']) {
        done[c - 'a'] = 1;
        cnt[c - 'a']++;
      }
    }

    for (int j = 0; j < s.size(); j++) {
      if (j)
        w[j] = w[j - 1];
      w[j][s[j] - 'a']--;
    }
    insert(root, s, 0, w);
  }

  ll ans = 0;
  fo(i, 0, n) {
    if (v[i].size() > 1) {
      ans += qry1(root, v[i], 0) - 1;
      ans += qry2(root, v[i], 0, v[i].back() - 'a') - 1;
    } else
      ans += cnt[v[i][0] - 'a'] - 1;
  }
  cout << ans;
}