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

#include <bits/stdc++.h>
using namespace std;
#define rep(i, n) for (int i = 0; i < (int)(n); i++)
#define rep_range(i, s, e) for (int i = (s); i < (int)(e); i++)
#define rep_rev(i, s, e) for (int i = (s)-1; i >= (int)(e); i--)
#define sort_vec(v) sort((v).begin(), (v).end())
typedef int64_t lint;
// lint mod = 998244353;
// lint mod = 1000000007;
template <typename T1, typename T2> inline bool chmax(T1 &a, T2 b) {
  return a < b && (a = b, true);
}
template <typename T1, typename T2> inline bool chmin(T1 &a, T2 b) {
  return a > b && (a = b, true);
}

// https://ei1333.github.io/luzhiled/snippets/structure/trie.html
template <int char_size> struct TrieNode {
  int par;
  int depth;
  array<int, char_size> childs;
  int exist;
  vector<int> accept;

  array<int, char_size> hist;

  TrieNode(int par, int depth) : par(par), depth(depth), exist(0) {
    childs.fill(-1);
    hist = {};
  }
};

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(-1, 0)); }

  void add(const string &str, int id) {
    int depth = 0;
    int cur = root;
    while (depth < str.size()) {
      const int c = str[depth] - margin;
      if (nodes[cur].childs[c] == -1) {
        nodes[cur].childs[c] = (int)nodes.size();
        nodes.push_back(Node(cur, depth + 1));
      }
      cur = nodes[cur].childs[c];
      depth++;
    }
    nodes[cur].accept.push_back(id);

    array<int, char_size> itemhist = {};
    cur = nodes[cur].par; // add対象文字列は長さ1以上なのでok
    while (true) {
      const int c = str[nodes[cur].depth] - margin;

      itemhist[c] = 1;
      nodes[cur].exist++;
      rep(c, char_size) nodes[cur].hist[c] += itemhist[c];

      cur = nodes[cur].par;
      if (cur == -1)
        break;
    }
  }

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

  void query(const string &str, const function<void(array<int, char_size>)> &f,
             int str_index, int node_index) {
    if (str_index == str.size()) {
      f(nodes[node_index].hist);
      return;
    } else {
      const int c = str[str_index] - margin;
      if (nodes[node_index].childs[c] == -1)
        return;
      query(str, f, str_index + 1, nodes[node_index].childs[c]);
    }
  }

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

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

  int size() const { return ((int)nodes.size()); }
};

lint app() {
  Trie<26, 'a'> trie;

  int n;
  cin >> n;

  vector<string> s(n);
  rep(i, n) {
    string st;
    cin >> st;
    reverse(st.begin(), st.end());
    s[i] = st;
    trie.add(st);
  }

  lint ans = 0;
  rep(i, n) {
    string q = s[i].substr(0, s[i].size() - 1);
    char c = s[i].back();

    auto f = [&](array<int, 26> hist) {
      // rep(i,26) cout<<hist[i];
      // cout<<endl;
      // cout<<q<<' '<<c<<endl;
      ans += hist[c - 'a'] - 1;
    };
    trie.query(q, f);
  }
  return ans;
}

int main() {
  cout << app() << endl;
  return 0;
}


Error: WA

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
#define rep(i, n) for (int i = 0; i < (int)(n); i++)
#define rep_range(i, s, e) for (int i = (s); i < (int)(e); i++)
#define rep_rev(i, s, e) for (int i = (s)-1; i >= (int)(e); i--)
#define sort_vec(v) sort((v).begin(), (v).end())
typedef int64_t lint;
// lint mod = 998244353;
// lint mod = 1000000007;
template <typename T1, typename T2> inline bool chmax(T1 &a, T2 b) {
  return a < b && (a = b, true);
}
template <typename T1, typename T2> inline bool chmin(T1 &a, T2 b) {
  return a > b && (a = b, true);
}

// https://ei1333.github.io/luzhiled/snippets/structure/trie.html
template <int char_size> struct TrieNode {
  int nxt[char_size];

  int exist;
  vector<int> accept;

  array<int, char_size> hist;

  TrieNode() : exist(0) {
    memset(nxt, -1, sizeof(nxt));
    rep(c, char_size) hist[c] = 0;
  }
};

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,
                    array<int, char_size> childhist) {
    ++nodes[node].exist;
    rep(c, char_size) nodes[node].hist[c] += childhist[c];
  }

  array<int, char_size> add(const string &str, int str_index, int node_index,
                            int id) {
    if (str_index == str.size()) {
      update_direct(node_index, id);
      array<int, char_size> hist;
      rep(c, char_size) hist[c] = 0;
      return hist;
    } 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());
      }
      auto childhist = add(str, str_index + 1, nodes[node_index].nxt[c], id);
      childhist[c] = 1;
      update_child(node_index, nodes[node_index].nxt[c], id, childhist);
      return childhist;
    }
  }

  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(array<int, char_size>)> &f,
             int str_index, int node_index) {
    if (str_index == str.size()) {
      f(nodes[node_index].hist);
      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(array<int, char_size>)> &f) {
    query(str, f, 0, 0);
  }

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

  int size() const { return ((int)nodes.size()); }
};

lint app() {
  Trie<26, 'a'> trie;

  int n;
  cin >> n;

  vector<string> s(n);
  rep(i, n) {
    string st;
    cin >> st;
    reverse(st.begin(), st.end());
    s[i] = st;
    trie.add(st);
    if (i > 10000)
      break;
  }

  return -1;
}

int main() {
  cout << app() << endl;
  return 0;
}