Problem p03159 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <algorithm>
#include <bitset>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <ctime>
#include <deque>
#include <functional>
#include <iomanip>
#include <iostream>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <sstream>
#include <stack>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>
using namespace std;

#define REP(i, s) for (int i = 0; i < s; ++i)
#define ALL(v) (v.begin(), v.end())
#define COUT(x) cout << #x << " = " << (x) << " (L" << __LINE__ << ")" << endl
#define EACH(i, s)                                                             \
  for (__typeof__((s).begin()) i = (s).begin(); i != (s).end(); ++i)

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;
}
template <class T1, class T2> ostream &operator<<(ostream &s, pair<T1, T2> P) {
  return s << '<' << P.first << ", " << P.second << '>';
}
template <class T> ostream &operator<<(ostream &s, vector<T> P) {
  for (int i = 0; i < P.size(); ++i) {
    if (i > 0) {
      s << " ";
    }
    s << P[i];
  }
  return s;
}
template <class T> ostream &operator<<(ostream &s, vector<vector<T>> P) {
  for (int i = 0; i < P.size(); ++i) {
    s << endl << P[i];
  }
  return s << endl;
}
template <class T> ostream &operator<<(ostream &s, set<T> P) {
  EACH(it, P) { s << "<" << *it << "> "; }
  return s << endl;
}
template <class T1, class T2> ostream &operator<<(ostream &s, map<T1, T2> P) {
  EACH(it, P) { s << "<" << it->first << "->" << it->second << "> "; }
  return s << endl;
}

const long long INF = 1LL << 60;
const int MAX = 5100;

int N;
vector<long long> A;
vector<vector<int>> G;

int num[MAX];
long long dp[MAX][MAX][2]; // vの部分木、vを含めず何個か、(全部正かどうか)
long long sdp[2][MAX][2];

void rec(int v, int p) {
  int numc = 0;
  num[v] = 1;
  for (auto ch : G[v]) {
    if (ch == p)
      continue;
    rec(ch, v);
    ++numc;
    num[v] += num[ch];
  }

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

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

  int i = 0;
  for (auto ch : G[v]) {
    if (ch == p)
      continue;
    for (int j = 0; j <= num[v]; ++j) {
      for (int k = 0; k <= num[ch] && j + k <= num[v]; ++k) {
        chmin(sdp[(i + 1) % 2][j + k][0], sdp[i % 2][j][0] + dp[ch][k][0]);
        chmin(sdp[(i + 1) % 2][j + k][0], sdp[i % 2][j][0] + dp[ch][k][1]);
        if (dp[ch][k][0] < 0)
          chmin(sdp[(i + 1) % 2][j + k + 1][0], sdp[i % 2][j][0]);
        if (dp[ch][k][1] < INF / 2)
          chmin(sdp[(i + 1) % 2][j + k + 1][0], sdp[i % 2][j][0]);
        if (A[v] > 0) {
          chmin(sdp[(i + 1) % 2][j + k][1], sdp[i % 2][j][1] + dp[ch][k][1]);
          if (dp[ch][k][0] < 0)
            chmin(sdp[(i + 1) % 2][j + k + 1][1], sdp[i % 2][j][1]);
          if (dp[ch][k][1] < INF / 2)
            chmin(sdp[(i + 1) % 2][j + k + 1][1], sdp[i % 2][j][1]);
        }
      }
    }
    for (int j = 0; j <= num[v]; ++j)
      sdp[i % 2][j][0] = sdp[i % 2][j][1] = INF;
    ++i;
  }
  for (int j = 0; j <= num[v]; ++j) {
    dp[v][j][0] = sdp[numc % 2][j][0];
    dp[v][j][1] = sdp[numc % 2][j][1];

    // cout << v << ", " << j << ": " << make_pair(dp[v][j][0], dp[v][j][1]) <<
    // endl;
  }
}

int main() {
  while (cin >> N) {
    A.resize(N);
    for (int i = 0; i < N; ++i)
      cin >> A[i];
    G.assign(N, vector<int>());
    for (int i = 0; i < N - 1; ++i) {
      int u, v;
      cin >> u >> v;
      --u, --v;
      G[u].push_back(v);
      G[v].push_back(u);
    }
    for (int i = 0; i < MAX; ++i)
      for (int j = 0; j < MAX; ++j)
        dp[i][j][0] = dp[i][j][1] = INF;
    rec(0, -1);
    int res = N;
    for (int i = 0; i <= N; ++i) {
      if (dp[0][i][0] < 0)
        chmin(res, i);
      if (dp[0][i][1] < INF / 2)
        chmin(res, i);
    }
    cout << res << endl;
  }
}


Fixed submission:
#include <algorithm>
#include <bitset>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <ctime>
#include <deque>
#include <functional>
#include <iomanip>
#include <iostream>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <sstream>
#include <stack>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>
using namespace std;

#define REP(i, s) for (int i = 0; i < s; ++i)
#define ALL(v) (v.begin(), v.end())
#define COUT(x) cout << #x << " = " << (x) << " (L" << __LINE__ << ")" << endl
#define EACH(i, s)                                                             \
  for (__typeof__((s).begin()) i = (s).begin(); i != (s).end(); ++i)

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;
}
template <class T1, class T2> ostream &operator<<(ostream &s, pair<T1, T2> P) {
  return s << '<' << P.first << ", " << P.second << '>';
}
template <class T> ostream &operator<<(ostream &s, vector<T> P) {
  for (int i = 0; i < P.size(); ++i) {
    if (i > 0) {
      s << " ";
    }
    s << P[i];
  }
  return s;
}
template <class T> ostream &operator<<(ostream &s, vector<vector<T>> P) {
  for (int i = 0; i < P.size(); ++i) {
    s << endl << P[i];
  }
  return s << endl;
}
template <class T> ostream &operator<<(ostream &s, set<T> P) {
  EACH(it, P) { s << "<" << *it << "> "; }
  return s << endl;
}
template <class T1, class T2> ostream &operator<<(ostream &s, map<T1, T2> P) {
  EACH(it, P) { s << "<" << it->first << "->" << it->second << "> "; }
  return s << endl;
}

const long long INF = 1LL << 60;
const int MAX = 5100;

int N;
vector<long long> A;
vector<vector<int>> G;

int num[MAX];
long long dp[MAX][MAX][2]; // vの部分木、vを含めず何個か、(全部正かどうか)
long long sdp[2][MAX][2];

void rec(int v, int p) {
  int numc = 0;
  num[v] = 1;
  for (auto ch : G[v]) {
    if (ch == p)
      continue;
    rec(ch, v);
    ++numc;
    num[v] += num[ch];
  }

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

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

  int i = 0;
  for (auto ch : G[v]) {
    if (ch == p)
      continue;
    for (int j = 0; j <= num[v]; ++j) {
      if (sdp[i % 2][j][0] > INF / 2)
        break;
      for (int k = 0; k <= num[ch] && j + k <= num[v]; ++k) {
        chmin(sdp[(i + 1) % 2][j + k][0], sdp[i % 2][j][0] + dp[ch][k][0]);
        chmin(sdp[(i + 1) % 2][j + k][0], sdp[i % 2][j][0] + dp[ch][k][1]);
        if (dp[ch][k][0] < 0)
          chmin(sdp[(i + 1) % 2][j + k + 1][0], sdp[i % 2][j][0]);
        if (dp[ch][k][1] < INF / 2)
          chmin(sdp[(i + 1) % 2][j + k + 1][0], sdp[i % 2][j][0]);
        if (A[v] > 0) {
          chmin(sdp[(i + 1) % 2][j + k][1], sdp[i % 2][j][1] + dp[ch][k][1]);
          if (dp[ch][k][0] < 0)
            chmin(sdp[(i + 1) % 2][j + k + 1][1], sdp[i % 2][j][1]);
          if (dp[ch][k][1] < INF / 2)
            chmin(sdp[(i + 1) % 2][j + k + 1][1], sdp[i % 2][j][1]);
        }
      }
    }
    for (int j = 0; j <= num[v]; ++j)
      sdp[i % 2][j][0] = sdp[i % 2][j][1] = INF;
    ++i;
  }
  for (int j = 0; j <= num[v]; ++j) {
    dp[v][j][0] = sdp[numc % 2][j][0];
    dp[v][j][1] = sdp[numc % 2][j][1];

    // cout << v << ", " << j << ": " << make_pair(dp[v][j][0], dp[v][j][1]) <<
    // endl;
  }
}

int main() {
  while (cin >> N) {
    A.resize(N);
    for (int i = 0; i < N; ++i)
      cin >> A[i];
    G.assign(N, vector<int>());
    for (int i = 0; i < N - 1; ++i) {
      int u, v;
      cin >> u >> v;
      --u, --v;
      G[u].push_back(v);
      G[v].push_back(u);
    }
    for (int i = 0; i < MAX; ++i)
      for (int j = 0; j < MAX; ++j)
        dp[i][j][0] = dp[i][j][1] = INF;
    rec(0, -1);
    int res = N;
    for (int i = 0; i <= N; ++i) {
      if (dp[0][i][0] < 0)
        chmin(res, i);
      if (dp[0][i][1] < INF / 2)
        chmin(res, i);
    }
    cout << res << endl;
  }
}