Problem p03159 (C++)

Buggy submission (status: Time Limit Exceeded):
/**
 *    author:  otera
 **/
#include <algorithm>
#include <bitset>
#include <cassert>
#include <ciso646>
#include <cmath>
#include <complex>
#include <cstdio>
#include <cstring>
#include <deque>
#include <functional>
#include <iomanip>
#include <iostream>
#include <map>
#include <queue>
#include <random>
#include <set>
#include <stack>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>
using namespace std;

#define int long long
typedef long long ll;
typedef unsigned long long ul;
typedef unsigned int ui;
typedef long double ld;
const int inf = 1e9 + 7;
const ll INF = 1LL << 60;
const ll mod = 1e9 + 7;
#define rep(i, n) for (int i = 0; i < n; i++)
#define per(i, n) for (int i = n - 1; i >= 0; i--)
#define Rep(i, sta, n) for (int i = sta; i < n; i++)
#define rep1(i, n) for (int i = 1; i <= n; i++)
#define per1(i, n) for (int i = n; i >= 1; i--)
#define Rep1(i, sta, n) for (int i = sta; i <= n; i++)
typedef complex<ld> Point;
const ld eps = 1e-8;
const ld pi = acos(-1.0);
typedef pair<int, int> P;
typedef pair<ld, ld> LDP;
typedef pair<ll, ll> LP;
#define fr first
#define sc second
#define all(c) c.begin(), c.end()
#define pb push_back
#define debug(x) cerr << #x << " = " << (x) << endl;
template <class T> inline bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return 1;
  }
  return 0;
}
template <class T> inline bool chmin(T &a, T b) {
  if (a > b) {
    a = b;
    return 1;
  }
  return 0;
}

const int MAX = 5100;

int n;
vector<int> a;
vector<vector<int>> g;

int num[MAX];
int dp[MAX][MAX][2];
int sdp[MAX][2];

void rec(int v, int p) {
  num[v] = 1;
  for (int nv : g[v]) {
    if (nv == p)
      continue;
    rec(nv, v);
    num[v] += num[nv];
  }

  for (int i = 0; i <= num[v]; ++i) {
    sdp[i][0] = sdp[i][1] = INF;
  }

  sdp[0][0] = a[v];
  if (a[v] > 0)
    sdp[0][1] = a[v];

  int cur = 0;
  for (int nv : g[v]) {
    if (nv == p)
      continue;
    for (int i = cur; i >= 0; --i) {
      int tmp0 = sdp[i][0], tmp1 = sdp[i][1];
      sdp[i][0] = INF, sdp[i][1] = INF;
      for (int j = 0; j <= num[nv]; ++j) {
        chmin(sdp[i + j][0], tmp0 + dp[nv][j][0]);
        if (dp[nv][j][0] < 0 || dp[nv][j][1] < INF)
          chmin(sdp[i + j + 1][0], tmp0);
        if (a[v] > 0) {
          chmin(sdp[i + j][1], tmp1 + dp[nv][j][1]);
          if (dp[nv][j][0] < 0 || dp[nv][j][1] < INF)
            chmin(sdp[i + j + 1][1], tmp1);
        }
      }
    }
    cur += num[v];
  }

  for (int i = 0; i <= num[v]; ++i) {
    dp[v][i][0] = sdp[i][0];
    dp[v][i][1] = sdp[i][1];
  }
}

void solve() {
  cin >> n;
  a.resize(n);
  rep(i, n) { cin >> a[i]; }
  g.assign(n, vector<int>());
  rep(i, n - 1) {
    int u, v;
    cin >> u >> v;
    --u, --v;
    g[u].pb(v);
    g[v].pb(u);
  }
  for (int i = 0; i < n; ++i) {
    for (int j = 0; j <= MAX; ++j) {
      dp[i][j][0] = dp[i][j][1] = INF;
    }
  }
  rec(0, -1);
  int ans = n;
  for (int j = n; j >= 0; --j) {
    if (dp[0][j][0] < 0)
      ans = j;
    if (dp[0][j][1] < INF)
      ans = j;
  }
  cout << ans << endl;
}

signed main() {
  ios::sync_with_stdio(false);
  cin.tie(0);
  // cout << fixed << setprecision(10);
  // int t; cin >> t; rep(i, t)solve();
  solve();
  return 0;
}

Fixed submission:
/**
 *    author:  otera
 **/
#include <algorithm>
#include <bitset>
#include <cassert>
#include <ciso646>
#include <cmath>
#include <complex>
#include <cstdio>
#include <cstring>
#include <deque>
#include <functional>
#include <iomanip>
#include <iostream>
#include <map>
#include <queue>
#include <random>
#include <set>
#include <stack>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>
using namespace std;

#define int long long
typedef long long ll;
typedef unsigned long long ul;
typedef unsigned int ui;
typedef long double ld;
const int inf = 1e9 + 7;
const ll INF = 1LL << 60;
const ll mod = 1e9 + 7;
#define rep(i, n) for (int i = 0; i < n; i++)
#define per(i, n) for (int i = n - 1; i >= 0; i--)
#define Rep(i, sta, n) for (int i = sta; i < n; i++)
#define rep1(i, n) for (int i = 1; i <= n; i++)
#define per1(i, n) for (int i = n; i >= 1; i--)
#define Rep1(i, sta, n) for (int i = sta; i <= n; i++)
typedef complex<ld> Point;
const ld eps = 1e-8;
const ld pi = acos(-1.0);
typedef pair<int, int> P;
typedef pair<ld, ld> LDP;
typedef pair<ll, ll> LP;
#define fr first
#define sc second
#define all(c) c.begin(), c.end()
#define pb push_back
#define debug(x) cerr << #x << " = " << (x) << endl;
template <class T> inline bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return 1;
  }
  return 0;
}
template <class T> inline bool chmin(T &a, T b) {
  if (a > b) {
    a = b;
    return 1;
  }
  return 0;
}

const int MAX = 5100;

int n;
vector<int> a;
vector<vector<int>> g;

int num[MAX];
int dp[MAX][MAX][2];
int sdp[MAX][2];

void rec(int v, int p) {
  num[v] = 1;
  for (int nv : g[v]) {
    if (nv == p)
      continue;
    rec(nv, v);
    num[v] += num[nv];
  }

  for (int i = 0; i <= num[v]; ++i) {
    sdp[i][0] = sdp[i][1] = INF;
  }

  sdp[0][0] = a[v];
  if (a[v] > 0)
    sdp[0][1] = a[v];

  int cur = 0;
  for (int nv : g[v]) {
    if (nv == p)
      continue;
    for (int i = cur; i >= 0; --i) {
      int tmp0 = sdp[i][0], tmp1 = sdp[i][1];
      sdp[i][0] = INF, sdp[i][1] = INF;
      for (int j = 0; j <= num[nv]; ++j) {
        chmin(sdp[i + j][0], tmp0 + dp[nv][j][0]);
        if (dp[nv][j][0] < 0 || dp[nv][j][1] < INF)
          chmin(sdp[i + j + 1][0], tmp0);
        if (a[v] > 0) {
          chmin(sdp[i + j][1], tmp1 + dp[nv][j][1]);
          if (dp[nv][j][0] < 0 || dp[nv][j][1] < INF)
            chmin(sdp[i + j + 1][1], tmp1);
        }
      }
    }
    cur += num[nv];
  }

  for (int i = 0; i <= num[v]; ++i) {
    dp[v][i][0] = sdp[i][0];
    dp[v][i][1] = sdp[i][1];
  }
}

void solve() {
  cin >> n;
  a.resize(n);
  rep(i, n) { cin >> a[i]; }
  g.assign(n, vector<int>());
  rep(i, n - 1) {
    int u, v;
    cin >> u >> v;
    --u, --v;
    g[u].pb(v);
    g[v].pb(u);
  }
  for (int i = 0; i < n; ++i) {
    for (int j = 0; j <= MAX; ++j) {
      dp[i][j][0] = dp[i][j][1] = INF;
    }
  }
  rec(0, -1);
  int ans = n;
  for (int j = n; j >= 0; --j) {
    if (dp[0][j][0] < 0)
      ans = j;
    if (dp[0][j][1] < INF)
      ans = j;
  }
  cout << ans << endl;
}

signed main() {
  ios::sync_with_stdio(false);
  cin.tie(0);
  // cout << fixed << setprecision(10);
  // int t; cin >> t; rep(i, t)solve();
  solve();
  return 0;
}