Problem p03159 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <bits/stdc++.h>
#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(a) sort(a.begin(), a.end())
#define REVERSE(a) reverse(a.begin(), a.end())
#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 int long long
#define INF 1000000000000000000
using namespace std;

#define ANS(f)                                                                 \
  if (f)                                                                       \
    cout << "YES" << endl;                                                     \
  else                                                                         \
    cout << "NO" << endl;

typedef vector<int> vec;
typedef vector<vec> mat;
typedef pair<int, int> Pii;

template <typename T> void readv(vector<T> &a) { REP(i, a.size()) cin >> a[i]; }
void readi(vector<int> &a) {
  REP(i, a.size()) {
    cin >> a[i];
    a[i]--;
  }
}
void debug(mat m) {
  REP(i, m.size()) {
    REP(j, m[i].size()) { cout << m[i][j] << ","; }
    cout << endl;
  }
}

struct edge {
  int to, cost;
};

class Graph {
public:
  int V;
  vector<vector<edge>> G;

  vec gen, par;
  mat chi, genlist;
  int depth;

  Graph(int V) : V(V) { G = vector<vector<edge>>(V, vector<edge>(0)); }

  void add_edge(int from, int to, int cost) {
    G[from].push_back(edge({to, cost}));
  }

  void add_edge2(int v1, int v2, int cost) {
    add_edge(v1, v2, cost);
    add_edge(v2, v1, cost);
  }

  void dfs(int v, int d) {
    gen[v] = d;
    depth = max(depth, d + 1);
    REP(k, G[v].size()) {
      if (G[v][k].to == par[v])
        continue;
      par[G[v][k].to] = v;
      chi[v].push_back(G[v][k].to);
      dfs(G[v][k].to, d + 1);
    }
  }

  mat makegenlist() {
    REP(v, V) genlist[gen[v]].push_back(v);
    return genlist;
  }

  void analyzeTree(int root) {
    gen = vec(V);
    par = vec(V);
    chi = mat(V, vec(0));
    depth = 0;
    par[root] = -1;
    dfs(root, 0);

    genlist = mat(depth, vec(0));
    makegenlist();
  }
};

signed main() {

  int N;
  cin >> N;
  vec A(N);
  readv(A);
  Graph T(N);
  int U, V;
  REP(i, N - 1) {
    cin >> U >> V;
    T.add_edge2(U - 1, V - 1, 0);
  }
  T.analyzeTree(0);

  vector<mat> dp(N);
  vec sz(N, 0);
  IREP(d, T.depth) {
    for (int v : T.genlist[d]) {

      vector<mat> dq(1 + T.chi[v].size());

      sz[v] = 1;
      dq[0] = mat(2, vec(1, INF));
      if (A[v] < 0)
        dq[0][1][0] = A[v];
      else
        dq[0][0][0] = A[v];
      int i0 = 1;
      for (int chi : T.chi[v]) {
        sz[v] += sz[chi];
        dq[i0] = mat(2, vec(sz[v], INF));
        // 切る
        int m = INF;
        REP(k, sz[chi]) {
          if (dp[chi][0][k] < INF || dp[chi][1][k] < 0) {
            m = k;
            break;
          }
        }
        if (m < INF) {
          REP(k, dq[i0 - 1][0].size()) {
            dq[i0][0][k + 1 + m] = dq[i0 - 1][0][k];
            dq[i0][1][k + 1 + m] = dq[i0 - 1][1][k];
          }
        }
        // 繫げる
        REP(k, sz[v]) {
          REP(l, min(sz[chi], k + 1)) {
            if (!(k - l >= 0 && k - l < dq[i0 - 1][0].size()))
              continue;
            if (dp[chi][0][l] < INF && dq[i0 - 1][0][k - l] < INF)
              dq[i0][0][k] =
                  min(dq[i0][0][k], dp[chi][0][l] + dq[i0 - 1][0][k - l]);
            if (dp[chi][1][l] < INF && dq[i0 - 1][1][k - l] < INF)
              dq[i0][1][k] =
                  min(dq[i0][1][k], dp[chi][1][l] + dq[i0 - 1][1][k - l]);
            if (dp[chi][0][l] < INF && dq[i0 - 1][1][k - l] < INF)
              dq[i0][1][k] =
                  min(dq[i0][1][k], dp[chi][0][l] + dq[i0 - 1][1][k - l]);
            if (dp[chi][1][l] < INF && dq[i0 - 1][0][k - l] < INF)
              dq[i0][1][k] =
                  min(dq[i0][1][k], dp[chi][1][l] + dq[i0 - 1][0][k - l]);
          }
        }
        i0++;
      }

      dp[v] = dq[T.chi[v].size()];
      /*
      cout << v + 1 << endl;
      REP(i, sz[v]) cout << dp[v][0][i] << ",";
      cout << endl;
      REP(i, sz[v]) cout << dp[v][1][i] << ",";
      cout << endl << endl;*/
    }
  }
  int ans = INF;
  REP(k, sz[0]) {
    if (dp[0][0][k] < INF || dp[0][1][k] < 0) {
      ans = k;
      break;
    }
  }
  cout << ans;

  return 0;
}

Fixed submission:
#include <bits/stdc++.h>
#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(a) sort(a.begin(), a.end())
#define REVERSE(a) reverse(a.begin(), a.end())
#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 int long long
#define INF 1000000000000000000
using namespace std;

#define ANS(f)                                                                 \
  if (f)                                                                       \
    cout << "YES" << endl;                                                     \
  else                                                                         \
    cout << "NO" << endl;

typedef vector<int> vec;
typedef vector<vec> mat;
typedef pair<int, int> Pii;

template <typename T> void readv(vector<T> &a) { REP(i, a.size()) cin >> a[i]; }
void readi(vector<int> &a) {
  REP(i, a.size()) {
    cin >> a[i];
    a[i]--;
  }
}
void debug(mat m) {
  REP(i, m.size()) {
    REP(j, m[i].size()) { cout << m[i][j] << ","; }
    cout << endl;
  }
}

struct edge {
  int to, cost;
};

class Graph {
public:
  int V;
  vector<vector<edge>> G;

  vec gen, par;
  mat chi, genlist;
  int depth;

  Graph(int V) : V(V) { G = vector<vector<edge>>(V, vector<edge>(0)); }

  void add_edge(int from, int to, int cost) {
    G[from].push_back(edge({to, cost}));
  }

  void add_edge2(int v1, int v2, int cost) {
    add_edge(v1, v2, cost);
    add_edge(v2, v1, cost);
  }

  void dfs(int v, int d) {
    gen[v] = d;
    depth = max(depth, d + 1);
    REP(k, G[v].size()) {
      if (G[v][k].to == par[v])
        continue;
      par[G[v][k].to] = v;
      chi[v].push_back(G[v][k].to);
      dfs(G[v][k].to, d + 1);
    }
  }

  mat makegenlist() {
    REP(v, V) genlist[gen[v]].push_back(v);
    return genlist;
  }

  void analyzeTree(int root) {
    gen = vec(V);
    par = vec(V);
    chi = mat(V, vec(0));
    depth = 0;
    par[root] = -1;
    dfs(root, 0);

    genlist = mat(depth, vec(0));
    makegenlist();
  }
};

signed main() {

  int N;
  cin >> N;
  vec A(N);
  readv(A);
  Graph T(N);
  int U, V;
  REP(i, N - 1) {
    cin >> U >> V;
    T.add_edge2(U - 1, V - 1, 0);
  }
  T.analyzeTree(0);

  vector<mat> dp(N);
  vec sz(N, 0);
  IREP(d, T.depth) {
    for (int v : T.genlist[d]) {

      vector<mat> dq(1 + T.chi[v].size());

      sz[v] = 1;
      dq[0] = mat(2, vec(1, INF));
      if (A[v] < 0)
        dq[0][1][0] = A[v];
      else
        dq[0][0][0] = A[v];
      int i0 = 1;
      for (int chi : T.chi[v]) {
        sz[v] += sz[chi];
        dq[i0] = mat(2, vec(sz[v], INF));
        // 切る
        int m = INF;
        REP(k, sz[chi]) {
          if (dp[chi][0][k] < INF || dp[chi][1][k] < 0) {
            m = k;
            break;
          }
        }
        if (m < INF) {
          REP(k, dq[i0 - 1][0].size()) {
            dq[i0][0][k + 1 + m] = dq[i0 - 1][0][k];
            dq[i0][1][k + 1 + m] = dq[i0 - 1][1][k];
          }
        }
        // 繫げる
        REP(k, sz[chi]) {
          REP(l, dq[i0 - 1][0].size()) {
            if (dp[chi][0][k] < INF && dq[i0 - 1][0][l] < INF)
              dq[i0][0][k + l] =
                  min(dq[i0][0][k + l], dp[chi][0][k] + dq[i0 - 1][0][l]);
            if (dp[chi][1][k] < INF && dq[i0 - 1][1][l] < INF)
              dq[i0][1][k + l] =
                  min(dq[i0][1][k + l], dp[chi][1][k] + dq[i0 - 1][1][l]);
            if (dp[chi][0][k] < INF && dq[i0 - 1][1][l] < INF)
              dq[i0][1][k + l] =
                  min(dq[i0][1][k + l], dp[chi][0][k] + dq[i0 - 1][1][l]);
            if (dp[chi][1][k] < INF && dq[i0 - 1][0][l] < INF)
              dq[i0][1][k + l] =
                  min(dq[i0][1][k + l], dp[chi][1][k] + dq[i0 - 1][0][l]);
          }
        }
        i0++;
      }

      dp[v] = dq[T.chi[v].size()];
      /*
      cout << v + 1 << endl;
      REP(i, sz[v]) cout << dp[v][0][i] << ",";
      cout << endl;
      REP(i, sz[v]) cout << dp[v][1][i] << ",";
      cout << endl << endl;*/
    }
  }
  int ans = INF;
  REP(k, sz[0]) {
    if (dp[0][0][k] < INF || dp[0][1][k] < 0) {
      ans = k;
      break;
    }
  }
  cout << ans;

  return 0;
}