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

#include <bits/stdc++.h>
using namespace std;
/*#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
using namespace __gnu_pbds;
template<typename T> using gpp_set = tree<T, null_type, less<T>, rb_tree_tag,
tree_order_statistics_node_update>; template<typename T, typename L> using
gpp_map = tree<T, L, less<T>, rb_tree_tag, tree_order_statistics_node_update>;
template<typename T> using gpp_multiset = tree<T, null_type, less_equal<T>,
rb_tree_tag, tree_order_statistics_node_update>;*/
struct fast_ios {
  fast_ios() {
    cin.tie(0);
    ios::sync_with_stdio(false);
    cout << fixed << setprecision(20);
  };
} fast_ios_;
#define FOR(i, begin, end) for (int i = (begin); i < (end); i++)
#define REP(i, n) FOR(i, 0, n)
#define IFOR(i, begin, end) for (int i = (end)-1; i >= (begin); i--)
#define IREP(i, n) IFOR(i, 0, n)
#define Sort(v) sort(v.begin(), v.end())
#define Reverse(v) reverse(v.begin(), v.end())
#define all(v) v.begin(), v.end()
#define SZ(v) ((int)v.size())
#define Lower_bound(v, x)                                                      \\
  distance(v.begin(), lower_bound(v.begin(), v.end(), x))
#define Upper_bound(v, x)                                                      \\
  distance(v.begin(), upper_bound(v.begin(), v.end(), x))
#define Max(a, b) a = max(a, b)
#define Min(a, b) a = min(a, b)
#define bit(n) (1LL << (n))
#define bit_exist(x, n) ((x >> n) & 1)
#define debug(x) cout << #x << "=" << x << endl;
#define vdebug(v)                                                              \\
  {                                                                            \\
    cout << #v << "=" << endl;                                                 \\
    REP(i_debug, v.size()) { cout << v[i_debug] << ","; }                      \\
    cout << endl;                                                              \\
  }
#define mdebug(m)                                                              \\
  {                                                                            \\
    cout << #m << "=" << endl;                                                 \\
    REP(i_debug, m.size()) {                                                   \\
      REP(j_debug, m[i_debug].size()) { cout << m[i_debug][j_debug] << ","; }  \\
      cout << endl;                                                            \\
    }                                                                          \\
  }
#define Return(ans)                                                            \\
  {                                                                            \\
    cout << (ans) << endl;                                                     \\
    return 0;                                                                  \\
  }
#define pb push_back
#define f first
#define s second
#define int long long
#define INF 1000000000000000000
template <typename T> istream &operator>>(istream &is, vector<T> &v) {
  for (auto &x : v)
    is >> x;
  return is;
}
template <typename T> ostream &operator<<(ostream &os, vector<T> &v) {
  for (int i = 0; i < v.size(); i++) {
    cout << v[i];
    if (i != v.size() - 1)
      cout << endl;
  };
  return os;
}
template <typename T1, typename T2>
ostream &operator<<(ostream &os, pair<T1, T2> p) {
  cout << '(' << p.first << ',' << p.second << ')';
  return os;
}
template <typename T> void Out(T x) { cout << x << endl; }
template <typename T1, typename T2> void Ans(bool f, T1 y, T2 n) {
  if (f)
    Out(y);
  else
    Out(n);
}

using vec = vector<int>;
using mat = vector<vec>;
using Pii = pair<int, int>;
using PiP = pair<int, Pii>;
using PPi = pair<Pii, int>;
using bools = vector<bool>;
using pairs = vector<Pii>;

// int dx[4] = {1,0,-1,0};
// int dy[4] = {0,1,0,-1};
// char d[4] = {'D','R','U','L'};

// const int mod = 1000000007;
const int mod = 998244353;
// #define Add(x, y) x = (x + (y)) % mod
// #define Mult(x, y) x = (x * (y)) % mod

template <long long MOD> struct ModInt {

  using ll = long long;
  ll val;

  void setval(ll v) { val = v % MOD; };
  ModInt() : val(0) {}
  ModInt(ll v) { setval(v); };

  ModInt operator+(const ModInt &x) const { return ModInt(val + x.val); }
  ModInt operator-(const ModInt &x) const { return ModInt(val - x.val + MOD); }
  ModInt operator*(const ModInt &x) const { return ModInt(val * x.val); }
  ModInt operator/(const ModInt &x) const { return *this * x.inv(); }
  ModInt operator-() const { return ModInt(MOD - val); }
  ModInt operator+=(const ModInt &x) { return *this = *this + x; }
  ModInt operator-=(const ModInt &x) { return *this = *this - x; }
  ModInt operator*=(const ModInt &x) { return *this = *this * x; }
  ModInt operator/=(const ModInt &x) { return *this = *this / x; }

  friend ostream &operator<<(ostream &os, const ModInt &x) {
    os << x.val;
    return os;
  }
  friend istream &operator>>(istream &is, ModInt &x) {
    is >> x.val;
    x.val = (x.val % MOD + MOD) % MOD;
    return is;
  }

  ModInt pow(ll n) const {
    ModInt a = 1;
    if (n == 0)
      return a;
    int i0 = 64 - __builtin_clzll(n);
    for (int i = i0 - 1; i >= 0; i--) {
      a = a * a;
      if ((n >> i) & 1)
        a *= (*this);
    }
    return a;
  }
  ModInt inv() const { return this->pow(MOD - 2); }
};

using mint = ModInt<mod>;
mint pow(mint x, long long n) { return x.pow(n); }
// using mint = double; //for debug
using mvec = vector<mint>;
using mmat = vector<mvec>;

struct Combination {

  vector<mint> fact, invfact;

  Combination(int N) {
    fact = vector<mint>({mint(1)});
    invfact = vector<mint>({mint(1)});
    fact_initialize(N);
  }

  void fact_initialize(int N) {
    int i0 = fact.size();
    if (i0 >= N + 1)
      return;
    fact.resize(N + 1);
    invfact.resize(N + 1);
    for (int i = i0; i <= N; i++)
      fact[i] = fact[i - 1] * i;
    invfact[N] = (mint)1 / fact[N];
    for (int i = N - 1; i >= i0; i--)
      invfact[i] = invfact[i + 1] * (i + 1);
  }

  mint nCr(int n, int r) {
    if (n < 0 || r < 0 || r > n)
      return mint(0);
    if (fact.size() < n + 1)
      fact_initialize(n);
    return fact[n] * invfact[r] * invfact[n - r];
  }

  mint nPr(int n, int r) {
    if (n < 0 || r < 0 || r > n)
      return mint(0);
    if (fact.size() < n + 1)
      fact_initialize(n);
    return fact[n] * invfact[n - r];
  }
};

struct edge {
  int to, cost, id;
};

class Graph {
public:
  int N;
  vector<vector<edge>> G;
  bools f;
  mmat dp;

  Graph(int N) : N(N) {
    G = vector<vector<edge>>(N, vector<edge>(0));
    f = bools(N, false);
    dp = mmat(N, mvec(3, 0));
  }

  void add_Directed_edge(int from, int to, int cost = 1, int id = 0) {
    G[from].push_back(edge({to, cost, id}));
  }

  void add_Undirected_edge(int v1, int v2, int cost = 1, int id = 0) {
    add_Directed_edge(v1, v2, cost, id);
    add_Directed_edge(v2, v1, cost, id);
  }

  // calculate diameter of a Tree
  void dfs_for_calcdist(int v, int p, int c, vec &dist) {
    if (p != -1)
      dist[v] = dist[p] + c;
    for (auto e : G[v])
      if (e.to != p)
        dfs_for_calcdist(e.to, v, e.cost, dist);
  }

  int diameter(int &v1, int &v2, int s = 0) {
    vec dist(N, 0);
    dfs_for_calcdist(s, -1, -1, dist);
    v1 = distance(dist.begin(), max_element(dist.begin(), dist.end()));
    fill(dist.begin(), dist.end(), 0);
    dfs_for_calcdist(v1, -1, -1, dist);
    auto amax = max_element(dist.begin(), dist.end());
    v2 = distance(dist.begin(), amax);
    return *amax;
  }

  void dfs(int v, int p) {

    if (f[v])
      dp[v][1] = 1;
    else
      dp[v][0] = 1;

    if (p != -1) {
      for (auto e : G[v])
        if (e.to != p) {
          dfs(e.to, v);
          mvec tmp(3, 0);
          tmp[0] =
              dp[v][0] * (dp[e.to][0] * 3 + dp[e.to][1] * 2 + dp[e.to][2] * 2);
          tmp[1] =
              dp[v][0] * dp[e.to][1] +
              dp[v][1] * (dp[e.to][0] * 3 + dp[e.to][1] * 2 + dp[e.to][2] * 2);
          tmp[2] = dp[v][1] * dp[e.to][1] +
                   dp[v][2] * (dp[e.to][0] + dp[e.to][1] + dp[e.to][2]) * 3 +
                   (dp[v][0] + dp[v][1]) * dp[e.to][2];
          dp[v] = tmp;
        }
    } else {
      dp[v] = mvec(4, 0);
      dp[v][0] = 1;
      for (auto e : G[v])
        if (e.to != p) {
          dfs(e.to, v);
          mvec tmp(4, 0);
          REP(j, 4) {
            tmp[j] += dp[v][j] * (dp[e.to][0] * 3 + dp[e.to][1] + dp[e.to][2]);
            if ((j & 1) == 0)
              tmp[j | 1] += dp[v][j] * dp[e.to][1];
            if ((j & 2) == 0)
              tmp[j | 2] += dp[v][j] * dp[e.to][1];
          }
          dp[v] = tmp;
        }
    }
  }
};

signed main() {

  int N;
  cin >> N;
  Graph G(N);
  REP(i, N - 1) {
    int a, b;
    cin >> a >> b;
    G.add_Undirected_edge(a - 1, b - 1);
  }

  int v1, v2;
  int R = G.diameter(v1, v2);
  vec d1(N, 0), d2(N, 0);
  G.dfs_for_calcdist(v1, -1, -1, d1);
  G.dfs_for_calcdist(v2, -1, -1, d2);

  vec center;
  REP(i, N) if (d1[i] + d2[i] == R && llabs(d1[i] - d2[i]) <= 1) center.pb(i);

  fill(all(d1), 0);
  fill(all(d2), 0);
  mint ans;
  if (SZ(center) == 1) {
    G.dfs_for_calcdist(center[0], -1, -1, d1);
    REP(i, N) if (d1[i] == R / 2) G.f[i] = true;

    G.dfs(center[0], -1);
    ans = G.dp[center[0]][3] / 2;
  } else {
    G.dfs_for_calcdist(center[0], -1, -1, d1);
    G.dfs_for_calcdist(center[1], -1, -1, d2);
    REP(i, N) if (min(d1[i], d2[i]) == R / 2) G.f[i] = true;

    G.dfs(center[0], center[1]);
    G.dfs(center[1], center[0]);
    ans = G.dp[center[0]][1] * G.dp[center[1]][1];
  }
  Out(ans);

  return 0;
}


Error: -6
70b925e6-ecc6-45b1-b43a-8f0bef8fb082.out: /home/alex/Documents/bug-detection/input/Project_CodeNet/data/p02752/C++/s648581634.cpp:233: int main(): Assertion `false' failed.


Fixed submission:

#include <bits/stdc++.h>
using namespace std;
/*#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
using namespace __gnu_pbds;
template<typename T> using gpp_set = tree<T, null_type, less<T>, rb_tree_tag,
tree_order_statistics_node_update>; template<typename T, typename L> using
gpp_map = tree<T, L, less<T>, rb_tree_tag, tree_order_statistics_node_update>;
template<typename T> using gpp_multiset = tree<T, null_type, less_equal<T>,
rb_tree_tag, tree_order_statistics_node_update>;*/
struct fast_ios {
  fast_ios() {
    cin.tie(0);
    ios::sync_with_stdio(false);
    cout << fixed << setprecision(20);
  };
} fast_ios_;
#define FOR(i, begin, end) for (int i = (begin); i < (end); i++)
#define REP(i, n) FOR(i, 0, n)
#define IFOR(i, begin, end) for (int i = (end)-1; i >= (begin); i--)
#define IREP(i, n) IFOR(i, 0, n)
#define Sort(v) sort(v.begin(), v.end())
#define Reverse(v) reverse(v.begin(), v.end())
#define all(v) v.begin(), v.end()
#define SZ(v) ((int)v.size())
#define Lower_bound(v, x)                                                      \\
  distance(v.begin(), lower_bound(v.begin(), v.end(), x))
#define Upper_bound(v, x)                                                      \\
  distance(v.begin(), upper_bound(v.begin(), v.end(), x))
#define Max(a, b) a = max(a, b)
#define Min(a, b) a = min(a, b)
#define bit(n) (1LL << (n))
#define bit_exist(x, n) ((x >> n) & 1)
#define debug(x) cout << #x << "=" << x << endl;
#define vdebug(v)                                                              \\
  {                                                                            \\
    cout << #v << "=" << endl;                                                 \\
    REP(i_debug, v.size()) { cout << v[i_debug] << ","; }                      \\
    cout << endl;                                                              \\
  }
#define mdebug(m)                                                              \\
  {                                                                            \\
    cout << #m << "=" << endl;                                                 \\
    REP(i_debug, m.size()) {                                                   \\
      REP(j_debug, m[i_debug].size()) { cout << m[i_debug][j_debug] << ","; }  \\
      cout << endl;                                                            \\
    }                                                                          \\
  }
#define Return(ans)                                                            \\
  {                                                                            \\
    cout << (ans) << endl;                                                     \\
    return 0;                                                                  \\
  }
#define pb push_back
#define f first
#define s second
#define int long long
#define INF 1000000000000000000
template <typename T> istream &operator>>(istream &is, vector<T> &v) {
  for (auto &x : v)
    is >> x;
  return is;
}
template <typename T> ostream &operator<<(ostream &os, vector<T> &v) {
  for (int i = 0; i < v.size(); i++) {
    cout << v[i];
    if (i != v.size() - 1)
      cout << endl;
  };
  return os;
}
template <typename T1, typename T2>
ostream &operator<<(ostream &os, pair<T1, T2> p) {
  cout << '(' << p.first << ',' << p.second << ')';
  return os;
}
template <typename T> void Out(T x) { cout << x << endl; }
template <typename T1, typename T2> void Ans(bool f, T1 y, T2 n) {
  if (f)
    Out(y);
  else
    Out(n);
}

using vec = vector<int>;
using mat = vector<vec>;
using Pii = pair<int, int>;
using PiP = pair<int, Pii>;
using PPi = pair<Pii, int>;
using bools = vector<bool>;
using pairs = vector<Pii>;

// int dx[4] = {1,0,-1,0};
// int dy[4] = {0,1,0,-1};
// char d[4] = {'D','R','U','L'};

// const int mod = 1000000007;
const int mod = 998244353;
// #define Add(x, y) x = (x + (y)) % mod
// #define Mult(x, y) x = (x * (y)) % mod

template <long long MOD> struct ModInt {

  using ll = long long;
  ll val;

  void setval(ll v) { val = v % MOD; };
  ModInt() : val(0) {}
  ModInt(ll v) { setval(v); };

  ModInt operator+(const ModInt &x) const { return ModInt(val + x.val); }
  ModInt operator-(const ModInt &x) const { return ModInt(val - x.val + MOD); }
  ModInt operator*(const ModInt &x) const { return ModInt(val * x.val); }
  ModInt operator/(const ModInt &x) const { return *this * x.inv(); }
  ModInt operator-() const { return ModInt(MOD - val); }
  ModInt operator+=(const ModInt &x) { return *this = *this + x; }
  ModInt operator-=(const ModInt &x) { return *this = *this - x; }
  ModInt operator*=(const ModInt &x) { return *this = *this * x; }
  ModInt operator/=(const ModInt &x) { return *this = *this / x; }

  friend ostream &operator<<(ostream &os, const ModInt &x) {
    os << x.val;
    return os;
  }
  friend istream &operator>>(istream &is, ModInt &x) {
    is >> x.val;
    x.val = (x.val % MOD + MOD) % MOD;
    return is;
  }

  ModInt pow(ll n) const {
    ModInt a = 1;
    if (n == 0)
      return a;
    int i0 = 64 - __builtin_clzll(n);
    for (int i = i0 - 1; i >= 0; i--) {
      a = a * a;
      if ((n >> i) & 1)
        a *= (*this);
    }
    return a;
  }
  ModInt inv() const { return this->pow(MOD - 2); }
};

using mint = ModInt<mod>;
mint pow(mint x, long long n) { return x.pow(n); }
// using mint = double; //for debug
using mvec = vector<mint>;
using mmat = vector<mvec>;

struct Combination {

  vector<mint> fact, invfact;

  Combination(int N) {
    fact = vector<mint>({mint(1)});
    invfact = vector<mint>({mint(1)});
    fact_initialize(N);
  }

  void fact_initialize(int N) {
    int i0 = fact.size();
    if (i0 >= N + 1)
      return;
    fact.resize(N + 1);
    invfact.resize(N + 1);
    for (int i = i0; i <= N; i++)
      fact[i] = fact[i - 1] * i;
    invfact[N] = (mint)1 / fact[N];
    for (int i = N - 1; i >= i0; i--)
      invfact[i] = invfact[i + 1] * (i + 1);
  }

  mint nCr(int n, int r) {
    if (n < 0 || r < 0 || r > n)
      return mint(0);
    if (fact.size() < n + 1)
      fact_initialize(n);
    return fact[n] * invfact[r] * invfact[n - r];
  }

  mint nPr(int n, int r) {
    if (n < 0 || r < 0 || r > n)
      return mint(0);
    if (fact.size() < n + 1)
      fact_initialize(n);
    return fact[n] * invfact[n - r];
  }
};

struct edge {
  int to, cost, id;
};

class Graph {
public:
  int N;
  vector<vector<edge>> G;
  bools f;
  mmat dp;

  Graph(int N) : N(N) {
    G = vector<vector<edge>>(N, vector<edge>(0));
    f = bools(N, false);
    dp = mmat(N, mvec(3, 0));
  }

  void add_Directed_edge(int from, int to, int cost = 1, int id = 0) {
    G[from].push_back(edge({to, cost, id}));
  }

  void add_Undirected_edge(int v1, int v2, int cost = 1, int id = 0) {
    add_Directed_edge(v1, v2, cost, id);
    add_Directed_edge(v2, v1, cost, id);
  }

  // calculate diameter of a Tree
  void dfs_for_calcdist(int v, int p, int c, vec &dist) {
    if (p != -1)
      dist[v] = dist[p] + c;
    for (auto e : G[v])
      if (e.to != p)
        dfs_for_calcdist(e.to, v, e.cost, dist);
  }

  int diameter(int &v1, int &v2, int s = 0) {
    vec dist(N, 0);
    dfs_for_calcdist(s, -1, -1, dist);
    v1 = distance(dist.begin(), max_element(dist.begin(), dist.end()));
    fill(dist.begin(), dist.end(), 0);
    dfs_for_calcdist(v1, -1, -1, dist);
    auto amax = max_element(dist.begin(), dist.end());
    v2 = distance(dist.begin(), amax);
    return *amax;
  }

  void dfs(int v, int p) {

    if (f[v])
      dp[v][1] = 1;
    else
      dp[v][0] = 1;

    if (p != -1) {
      for (auto e : G[v])
        if (e.to != p) {
          dfs(e.to, v);
          mvec tmp(3, 0);
          tmp[0] =
              dp[v][0] * (dp[e.to][0] * 3 + dp[e.to][1] * 2 + dp[e.to][2] * 2);
          tmp[1] =
              dp[v][0] * dp[e.to][1] +
              dp[v][1] * (dp[e.to][0] * 3 + dp[e.to][1] * 2 + dp[e.to][2] * 2);
          tmp[2] = dp[v][1] * dp[e.to][1] +
                   dp[v][2] * (dp[e.to][0] + dp[e.to][1] + dp[e.to][2]) * 3;
          dp[v] = tmp;
        }
    } else {
      dp[v] = mvec(4, 0);
      dp[v][0] = 1;
      for (auto e : G[v])
        if (e.to != p) {
          dfs(e.to, v);
          mvec tmp(4, 0);
          REP(j, 4) {
            tmp[j] += dp[v][j] * (dp[e.to][0] * 3 + dp[e.to][1] + dp[e.to][2]);
            if ((j & 1) == 0)
              tmp[j | 1] += dp[v][j] * dp[e.to][1];
            if ((j & 2) == 0)
              tmp[j | 2] += dp[v][j] * dp[e.to][1];
          }
          dp[v] = tmp;
        }
    }
  }
};

signed main() {

  int N;
  cin >> N;
  Graph G(N);
  REP(i, N - 1) {
    int a, b;
    cin >> a >> b;
    G.add_Undirected_edge(a - 1, b - 1);
  }

  int v1, v2;
  int R = G.diameter(v1, v2);
  vec d1(N, 0), d2(N, 0);
  G.dfs_for_calcdist(v1, -1, -1, d1);
  G.dfs_for_calcdist(v2, -1, -1, d2);

  vec center;
  REP(i, N) if (d1[i] + d2[i] == R && llabs(d1[i] - d2[i]) <= 1) center.pb(i);

  fill(all(d1), 0);
  fill(all(d2), 0);
  mint ans;
  if (SZ(center) == 1) {
    G.dfs_for_calcdist(center[0], -1, -1, d1);
    REP(i, N) if (d1[i] == R / 2) G.f[i] = true;

    G.dfs(center[0], -1);
    ans = G.dp[center[0]][3] / 2;
    assert(false);
  } else {
    G.dfs_for_calcdist(center[0], -1, -1, d1);
    G.dfs_for_calcdist(center[1], -1, -1, d2);
    REP(i, N) if (min(d1[i], d2[i]) == R / 2) G.f[i] = true;

    G.dfs(center[0], center[1]);
    G.dfs(center[1], center[0]);
    ans = G.dp[center[0]][1] * G.dp[center[1]][1];
  }
  Out(ans);

  return 0;
}