Problem p03159 (C++)

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

using namespace std;

typedef long long ll;

const ll MOD = 1e9 + 7;
const ll INF = 1e18;

signed main() {
  int n;
  cin >> n;
  vector<int> a(n);
  for (auto &y : a)
    cin >> y;
  vector<int> u(n - 1);
  vector<int> v(n - 1);
  for (int i = 0; i < n - 1; ++i) {
    cin >> u[i] >> v[i];
    --u[i];
    --v[i];
  }
  vector<vector<int>> edges(n);
  for (int i = 0; i < n - 1; ++i) {
    edges[u[i]].push_back(v[i]);
    edges[v[i]].push_back(u[i]);
  }
  vector<int> size(n, 1);

  bitset<5000> flag;
  function<void(int)> siz = [&](int x) {
    for (auto &y : edges[x])
      if (!flag[y]) {
        flag[y] = 1;
        siz(y);
        size[x] += size[y];
      }
  };
  flag[0] = 1;
  siz(0);
  flag.reset();

  function<vector<vector<ll>>(int)> f = [&](int x) {
    vector<vector<ll>> dp(1, vector<ll>(2, INF));
    dp[0][0] = a[x];
    if (a[x] > 0)
      dp[0][1] = a[x];

    int sizesum = 0;
    for (auto &y : edges[x]) {
      if (flag[y])
        continue;
      flag[y] = 1;
      auto from = f(y);

      vector<vector<ll>> next(sizesum + size[y] + 1, vector<ll>(2, INF));

      for (int i = 0; i <= sizesum; ++i) {
        for (int j = 0; j < size[y]; ++j) {
          next[i + j][0] = min(next[i + j][0], dp[i][0] + from[j][0]);
          next[i + j][0] = min(next[i + j][0], dp[i][0] + from[j][1]);
          next[i + j][0] = min(next[i + j][0], dp[i][1] + from[j][0]);
          next[i + j][0] = min(next[i + j][0], dp[i][1] + from[j][1]);
          if (a[x] > 0 && from[j][1] != INF)
            next[i + j][1] = min(next[i + j][1], dp[i][1] + from[j][1]);
          // 分断が可能なら
          if (from[j][0] < 0 || from[j][1] != INF) {
            next[i + j + 1][0] = min(next[i + j + 1][0], dp[i][0]);
            if (a[x] > 0)
              next[i + j + 1][1] = min(next[i + j + 1][1], dp[i][1]);
          }
        }
      }
      dp = move(next);
      sizesum += size[y];
    }
    return dp;
  };
  flag[0] = 1;
  auto dp = f(0);
  int ans = 1e9;
  for (int j = 0; j <= n; ++j) {
    if (dp[j][0] < 0)
      ans = min(ans, j);
    if (dp[j][1] != INF)
      ans = min(ans, j);
  }
  cout << ans << endl;
}


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

using namespace std;

typedef long long ll;

const ll MOD = 1e9 + 7;
const ll INF = 1e18;

signed main() {
  int n;
  cin >> n;
  vector<int> a(n);
  for (auto &y : a)
    cin >> y;
  vector<int> u(n - 1);
  vector<int> v(n - 1);
  for (int i = 0; i < n - 1; ++i) {
    cin >> u[i] >> v[i];
    --u[i];
    --v[i];
  }
  vector<vector<int>> edges(n);
  for (int i = 0; i < n - 1; ++i) {
    edges[u[i]].push_back(v[i]);
    edges[v[i]].push_back(u[i]);
  }
  vector<int> size(n, 1);

  bitset<5000> flag;
  function<void(int)> siz = [&](int x) {
    for (auto &y : edges[x])
      if (!flag[y]) {
        flag[y] = 1;
        siz(y);
        size[x] += size[y];
      }
  };
  flag[0] = 1;
  siz(0);
  flag.reset();

  function<vector<vector<ll>>(int)> f = [&](int x) {
    vector<vector<ll>> dp(1, vector<ll>(2, INF));
    dp[0][0] = a[x];
    if (a[x] > 0)
      dp[0][1] = a[x];

    int sizesum = 0;
    for (auto &y : edges[x]) {
      if (flag[y])
        continue;
      flag[y] = 1;
      auto from = f(y);

      vector<vector<ll>> next(sizesum + size[y] + 1, vector<ll>(2, INF));

      for (int i = 0; i <= sizesum; ++i) {
        for (int j = 0; j < size[y]; ++j) {
          next[i + j][0] = min(next[i + j][0], dp[i][0] + from[j][0]);
          next[i + j][0] = min(next[i + j][0], dp[i][0] + from[j][1]);
          next[i + j][0] = min(next[i + j][0], dp[i][1] + from[j][0]);
          next[i + j][0] = min(next[i + j][0], dp[i][1] + from[j][1]);
          if (a[x] > 0 && from[j][1] != INF)
            next[i + j][1] = min(next[i + j][1], dp[i][1] + from[j][1]);
          // 分断が可能なら
          if (from[j][0] < 0 || from[j][1] != INF) {
            next[i + j + 1][0] = min(next[i + j + 1][0], dp[i][0]);
            if (a[x] > 0)
              next[i + j + 1][1] = min(next[i + j + 1][1], dp[i][1]);
          }
        }
      }
      dp = move(next);
      sizesum += size[y];
    }
    return dp;
  };
  flag[0] = 1;
  auto dp = f(0);
  int ans = 1e9;
  for (int j = 0; j < (int)dp.size(); ++j) {
    if (dp[j][0] < 0)
      ans = min(ans, j);
    if (dp[j][1] != INF)
      ans = min(ans, j);
  }
  cout << ans << endl;
}