Problem p03809 (C++)

Buggy submission (status: Runtime Error):
#include <bits/stdc++.h>

#define FOR(i, a, b) for (ll i = (a); i < (ll)(b); i++)
#define REP(i, n) FOR(i, 0, n)
#define YYS(x, arr) for (auto &x : arr)
#define ALL(x) (x).begin(), (x).end()
#define SORT(x) sort((x).begin(), (x).end())
#define REVERSE(x) reverse((x).begin(), (x).end())
#define UNIQUE(x) (x).erase(unique(ALL((x))), (x).end())
#define PW(x) (1LL << (x))
#define SZ(x) ((ll)(x).size())
#define SHOW(x) cout << #x << " = " << x << endl
#define SHOWA(x, n)                                                            \
  for (int yui = 0; yui < n; yui++) {                                          \
    cout << x[yui] << " ";                                                     \
  }                                                                            \
  cout << endl

#define pb emplace_back
#define fi first
#define se second

using namespace std;

typedef long double ld;
typedef long long int ll;
typedef pair<int, int> pi;
typedef pair<ll, ll> pl;
typedef vector<int> vi;
typedef vector<ll> vl;
typedef vector<bool> vb;
typedef vector<ld> vd;
typedef vector<pi> vpi;
typedef vector<pl> vpl;
typedef vector<vpl> gr;
typedef vector<vl> ml;
typedef vector<vd> md;
typedef vector<vi> mi;

const ll INF = (ll)1e9 + 10;
const ll INFLL = (ll)1e18 + 10;
const ld EPS = 1e-12;
const ll MOD = 1e9 + 7;

template <class T> T &chmin(T &a, const T &b) { return a = min(a, b); }
template <class T> T &chmax(T &a, const T &b) { return a = max(a, b); }
template <class T> inline T sq(T a) { return a * a; }

ll in() {
  long long int x;
  scanf("%lld", &x);
  return x;
}
char yuyushiki[1000010];
string stin() {
  scanf("%s", yuyushiki);
  return yuyushiki;
}

// head

int n;
ll a[100010];

vi G[100010];

bool ans = true;

ll dfs(int x, int p) {
  ll cnt = 0;
  vl v(0);
  YYS(w, G[x]) {
    if (w != p) {
      ll res = dfs(w, x);
      v.pb(res);
      cnt += res;
    }
  }
  if (SZ(v) == 0) {
    return a[x];
  }
  ll ma = *max_element(ALL(v));
  ll able = cnt - ma;
  ll del = cnt - a[x];
  if (a[x] - del < 0 or del < 0 or able < del) {
    ans = false;
    return 0;
  }
  // cout << x+1 << " " << a[x] - del << endl;
  return a[x] - del;
}

int main() {

  n = in();
  REP(i, n) { a[i] = in(); }

  REP(i, n - 1) {
    int a = in() - 1;
    int b = in() - 1;
    G[a].pb(b);
    G[b].pb(a);
  }

  if (n == 2) {
    if (a[0] == a[1]) {
      puts("YES");
    } else {
      puts("NO");
    }
  }

  int root = -1;
  REP(i, n) {
    if (SZ(G[i]) >= 2) {
      root = i;
      break;
    }
  }

  assert(root != -1);

  ll res = dfs(root, -1);
  if (res != 0) {
    ans = false;
  }

  if (ans) {
    puts("YES");
  } else {
    puts("NO");
  }

  return 0;
}


Fixed submission:
#include <bits/stdc++.h>

#define FOR(i, a, b) for (ll i = (a); i < (ll)(b); i++)
#define REP(i, n) FOR(i, 0, n)
#define YYS(x, arr) for (auto &x : arr)
#define ALL(x) (x).begin(), (x).end()
#define SORT(x) sort((x).begin(), (x).end())
#define REVERSE(x) reverse((x).begin(), (x).end())
#define UNIQUE(x) (x).erase(unique(ALL((x))), (x).end())
#define PW(x) (1LL << (x))
#define SZ(x) ((ll)(x).size())
#define SHOW(x) cout << #x << " = " << x << endl
#define SHOWA(x, n)                                                            \
  for (int yui = 0; yui < n; yui++) {                                          \
    cout << x[yui] << " ";                                                     \
  }                                                                            \
  cout << endl

#define pb emplace_back
#define fi first
#define se second

using namespace std;

typedef long double ld;
typedef long long int ll;
typedef pair<int, int> pi;
typedef pair<ll, ll> pl;
typedef vector<int> vi;
typedef vector<ll> vl;
typedef vector<bool> vb;
typedef vector<ld> vd;
typedef vector<pi> vpi;
typedef vector<pl> vpl;
typedef vector<vpl> gr;
typedef vector<vl> ml;
typedef vector<vd> md;
typedef vector<vi> mi;

const ll INF = (ll)1e9 + 10;
const ll INFLL = (ll)1e18 + 10;
const ld EPS = 1e-12;
const ll MOD = 1e9 + 7;

template <class T> T &chmin(T &a, const T &b) { return a = min(a, b); }
template <class T> T &chmax(T &a, const T &b) { return a = max(a, b); }
template <class T> inline T sq(T a) { return a * a; }

ll in() {
  long long int x;
  scanf("%lld", &x);
  return x;
}
char yuyushiki[1000010];
string stin() {
  scanf("%s", yuyushiki);
  return yuyushiki;
}

// head

int n;
ll a[100010];

vi G[100010];

bool ans = true;

ll dfs(int x, int p) {
  ll cnt = 0;
  vl v(0);
  YYS(w, G[x]) {
    if (w != p) {
      ll res = dfs(w, x);
      v.pb(res);
      cnt += res;
    }
  }
  if (SZ(v) == 0) {
    return a[x];
  }
  ll ma = *max_element(ALL(v));
  ll able = cnt - ma;
  ll del = cnt - a[x];
  if (a[x] - del < 0 or del < 0 or able < del) {
    ans = false;
    return 0;
  }
  // cout << x+1 << " " << a[x] - del << endl;
  return a[x] - del;
}

int main() {

  n = in();
  REP(i, n) { a[i] = in(); }

  REP(i, n - 1) {
    int a = in() - 1;
    int b = in() - 1;
    G[a].pb(b);
    G[b].pb(a);
  }

  if (n == 2) {
    if (a[0] == a[1]) {
      puts("YES");
    } else {
      puts("NO");
    }
    return 0;
  }

  int root = -1;
  REP(i, n) {
    if (SZ(G[i]) >= 2) {
      root = i;
      break;
    }
  }

  assert(root != -1);

  ll res = dfs(root, -1);
  if (res != 0) {
    ans = false;
  }

  if (ans) {
    puts("YES");
  } else {
    puts("NO");
  }

  return 0;
}