Problem p03159 (C++)

Buggy submission (status: Runtime Error):
#include <iostream>
#include <vector>

using namespace std;

const long long MAXN = 3e3 + 5;

vector<long long> v1[MAXN];
long long dp[MAXN][MAXN][3];
long long temp[MAXN][3];
long long sz[MAXN];
long long arr[MAXN];
void dfs(long long curr, long long par) {
  for (long long x : v1[curr]) {
    if (x != par) {
      dfs(x, curr);
    }
  }
  long long sum = 1; // current subtree sz
  dp[curr][0][arr[curr] < 0] =
      arr[curr]; // 0 edges cut, is a computer = 1, not a computer  = 0
  for (long long x : v1[curr]) {
    if (x != par) {
      for (long long i = 0; i < sum; i++) {
        for (long long j = 0; j < sz[x]; j++) {
          for (long long k1 = 0; k1 < 2; k1++) {
            for (long long k2 = 0; k2 < 2; k2++) {
              // merge components together(all possible cuts)
              temp[i + j][k1 | k2] =
                  min(temp[i + j][k1 | k2], dp[curr][i][k1] + dp[x][j][k2]);
              if ((k2 == 0 && dp[x][j][k2] < 1e18) ||
                  (k2 == 1 && dp[x][j][k2] < 0)) {
                // cut edge between curr and x, cuts++, whether or not a
                // computer exists depends on the current component itself. You
                // can only do this if x's subtree has a value less than 0 when
                // it has a computer, or if x's subtree doesn't currently have a
                // computer.
                temp[i + j + 1][k1] = min(temp[i + j + 1][k1], dp[curr][i][k1]);
              }
            }
          }
        }
      }
      //(actually O(n))
      sum += sz[x];
      for (long long i = 0; i < sum; i++) {
        for (long long k = 0; k < 2; k++) {
          dp[curr][i][k] = temp[i][k];
          temp[i][k] = 1e18;
          // update all values to dp
        }
      }
    }
  }
  sz[curr] = sum;
}
int main() {
  long long n;
  cin >> n;
  for (long long i = 1; i <= n; i++) {
    cin >> arr[i];
  }
  for (long long i = 1; i < n; i++) {
    long long u, v;
    cin >> u >> v;
    v1[u].push_back(v);
    v1[v].push_back(u);
  }
  for (long long i = 1; i <= n; i++) {
    for (long long j = 0; j <= n; j++) {
      for (long long k = 0; k <= 1; k++) {
        dp[i][j][k] = 1e18;
        temp[j][k] = 1e18;
      }
    }
  }
  dfs(1, 1);
  long long ans = n - 1;
  for (long long i = 0; i < n; i++) {
    if (dp[1][i][0] < 1e18) {
      ans = min(ans, i);
    }
    if (dp[1][i][1] < 0) {
      ans = min(ans, i);
    }
  }
  cout << ans << endl;
}


Fixed submission:
#include <iostream>
#include <vector>

using namespace std;

const long long MAXN = 5e3 + 5;

vector<long long> v1[MAXN];
long long dp[MAXN][MAXN][3];
long long temp[MAXN][3];
long long sz[MAXN];
long long arr[MAXN];
void dfs(long long curr, long long par) {
  for (long long x : v1[curr]) {
    if (x != par) {
      dfs(x, curr);
    }
  }
  long long sum = 1; // current subtree sz
  dp[curr][0][arr[curr] < 0] =
      arr[curr]; // 0 edges cut, is a computer = 1, not a computer  = 0
  for (long long x : v1[curr]) {
    if (x != par) {
      for (long long i = 0; i < sum; i++) {
        for (long long j = 0; j < sz[x]; j++) {
          for (long long k1 = 0; k1 < 2; k1++) {
            for (long long k2 = 0; k2 < 2; k2++) {
              // merge components together(all possible cuts)
              temp[i + j][k1 | k2] =
                  min(temp[i + j][k1 | k2], dp[curr][i][k1] + dp[x][j][k2]);
              if ((k2 == 0 && dp[x][j][k2] < 1e18) ||
                  (k2 == 1 && dp[x][j][k2] < 0)) {
                // cut edge between curr and x, cuts++, whether or not a
                // computer exists depends on the current component itself. You
                // can only do this if x's subtree has a value less than 0 when
                // it has a computer, or if x's subtree doesn't currently have a
                // computer.
                temp[i + j + 1][k1] = min(temp[i + j + 1][k1], dp[curr][i][k1]);
              }
            }
          }
        }
      }
      //(actually O(n))
      sum += sz[x];
      for (long long i = 0; i < sum; i++) {
        for (long long k = 0; k < 2; k++) {
          dp[curr][i][k] = temp[i][k];
          temp[i][k] = 1e18;
          // update all values to dp
        }
      }
    }
  }
  sz[curr] = sum;
}
int main() {
  long long n;
  cin >> n;
  for (long long i = 1; i <= n; i++) {
    cin >> arr[i];
  }
  for (long long i = 1; i < n; i++) {
    long long u, v;
    cin >> u >> v;
    v1[u].push_back(v);
    v1[v].push_back(u);
  }
  for (long long i = 1; i <= n; i++) {
    for (long long j = 0; j <= n; j++) {
      for (long long k = 0; k <= 1; k++) {
        dp[i][j][k] = 1e18;
        temp[j][k] = 1e18;
      }
    }
  }
  dfs(1, 1);
  long long ans = n - 1;
  for (long long i = 0; i < n; i++) {
    if (dp[1][i][0] < 1e18) {
      ans = min(ans, i);
    }
    if (dp[1][i][1] < 0) {
      ans = min(ans, i);
    }
  }
  cout << ans << endl;
}