Problem p03809 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <bits/stdc++.h>
using namespace std;

typedef long long ll;
typedef vector<ll> vl;
typedef vector<vl> vvl;
typedef pair<ll, ll> pll;
typedef vector<bool> vb;
const ll oo = 0x3f3f3f3f3f3f3f3f;
const double eps = 1e-9;
#define sz(c) ll((c).size())
#define all(c) begin(c), end(c)
#define FOR(i, a, b) for (ll i = (a); i < (b); i++)
#define FORD(i, a, b) for (ll i = (b)-1; i >= (a); i--)
#define mp make_pair
#define mt make_tuple
#define pb push_back
#define eb emplace_back
#define xx first
#define yy second
#define has(c, i) ((c).find(i) != end(c))
#define DBGDO(X)                                                               \
  ({                                                                           \
    if (1)                                                                     \
      cerr << "DBGDO: " << (#X) << " = " << (X) << endl;                       \
  })

ll n;
ll a[100100];
vl adj[100100];

bool ok = true;

ll dfs(ll i, ll p = -1) {
  if (sz(adj[i]) == 1) {
    if (p != -1)
      return a[i];
    assert(false);
  } else {
    ll sum = 0;
    for (ll j : adj[i])
      if (j != p) {
        ll cur = dfs(j, i);
        if (cur > a[i])
          ok = false;
        sum += cur;
      }
    ll res = 2 * a[i] - sum;
    if (res > a[i])
      ok = false;
    if (res < 0)
      ok = false;
    return res;
  }
}

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(NULL);

  cin >> n;
  FOR(i, 0, n) cin >> a[i];
  FOR(i, 1, n) {
    ll u, v;
    cin >> u >> v;
    u--, v--;
    adj[u].pb(v), adj[v].pb(u);
  }

  if (n == 2) {
    ok = a[0] == a[1];
  } else {
    FOR(i, 0, n) if (sz(adj[i]) > 1) {
      ll res = dfs(i);
      if (res)
        ok = false;
    }
  }
  cout << (ok ? "YES" : "NO") << endl;
}


Fixed submission:
#include <bits/stdc++.h>
using namespace std;

typedef long long ll;
typedef vector<ll> vl;
typedef vector<vl> vvl;
typedef pair<ll, ll> pll;
typedef vector<bool> vb;
const ll oo = 0x3f3f3f3f3f3f3f3f;
const double eps = 1e-9;
#define sz(c) ll((c).size())
#define all(c) begin(c), end(c)
#define FOR(i, a, b) for (ll i = (a); i < (b); i++)
#define FORD(i, a, b) for (ll i = (b)-1; i >= (a); i--)
#define mp make_pair
#define mt make_tuple
#define pb push_back
#define eb emplace_back
#define xx first
#define yy second
#define has(c, i) ((c).find(i) != end(c))
#define DBGDO(X)                                                               \
  ({                                                                           \
    if (1)                                                                     \
      cerr << "DBGDO: " << (#X) << " = " << (X) << endl;                       \
  })

ll n;
ll a[100100];
vl adj[100100];

bool ok = true;

ll dfs(ll i, ll p = -1) {
  if (sz(adj[i]) == 1) {
    if (p != -1)
      return a[i];
    assert(false);
  } else {
    ll sum = 0;
    for (ll j : adj[i])
      if (j != p) {
        ll cur = dfs(j, i);
        if (cur > a[i])
          ok = false;
        sum += cur;
      }
    ll res = 2 * a[i] - sum;
    if (res > a[i])
      ok = false;
    if (res < 0)
      ok = false;
    return res;
  }
}

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(NULL);

  cin >> n;
  FOR(i, 0, n) cin >> a[i];
  FOR(i, 1, n) {
    ll u, v;
    cin >> u >> v;
    u--, v--;
    adj[u].pb(v), adj[v].pb(u);
  }

  if (n == 2) {
    ok = a[0] == a[1];
  } else {
    FOR(i, 0, n) if (sz(adj[i]) > 1) {
      ll res = dfs(i);
      if (res)
        ok = false;
      break;
    }
  }
  cout << (ok ? "YES" : "NO") << endl;
}