Problem p00384 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef double D;
typedef pair<ll, ll> P;
#define M 1000000007
#define F first
#define S second
#define PB push_back
#define INF 100000000000000000
#define MP make_pair
struct t {
  ll r, s1, s2;
  t(void) { r = 0, s1 = 0, s2 = 0; }
  void up(ll x) {
    if (s1 < x)
      swap(s1, x);
    s2 = max(s2, x);
  }
  void mn(ll x) {
    if (s1 == x) {
      s1 = s2;
      s2 = 0;
    }
  }
};
ll n, cs[100005], ans;
vector<ll> g[100005];
map<ll, t> dp[100005];
t dfs(ll v, ll p, ll f) {
  if (dp[v].count(p))
    return dp[v][p];
  t res;
  if (p == -1 || f) {
    for (ll u : g[v]) {
      if (u == p)
        continue;
      t x = dfs(u, v, f);
      x.r -= 2;
      x.s1--;
      if (x.r > 0) {
        res.r += x.r;
        res.up(-x.r + x.s1);
      } else
        res.up(x.s1);
    }
    res.r += cs[v];
    res.s1 += res.r;
    res.s2 += res.r;
  } else {
    res = dfs(v, -1, 0);
    t x = dfs(p, v, 0);
    x.r -= 2;
    x.s1--;
    if (x.r > 0) {
      res.r -= x.r;
      res.mn(res.r + x.s1);
      res.s1 -= x.r;
    } else
      res.mn(res.r + x.s1);
  }
  return dp[v][p] = res;
}
int main(void) {
  scanf("%lld", &n);
  ll ma = -INF;
  for (int i = 0; i < n; i++) {
    scanf("%lld", cs + i);
    ma = max(ma, cs[i]);
  }
  if (ma <= 0) {
    printf("%lld\n", ma);
    return 0;
  }
  for (int i = 0; i < n - 1; i++) {
    ll a, b;
    scanf("%lld%lld", &a, &b);
    g[--a].PB(--b);
    g[b].PB(a);
  }
  dfs(0, -1, 1);
  for (int i = 0; i < n; i++)
    ans = max(ans, dfs(i, -1, 0).s1);
  printf("%lld\n", ans);
}


Fixed submission:
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef double D;
typedef pair<ll, ll> P;
#define M 1000000007
#define F first
#define S second
#define PB push_back
#define INF 100000000000000000
#define MP make_pair
struct t {
  ll r, s1, s2;
  t(void) { r = 0, s1 = 0, s2 = 0; }
  void up(ll x) {
    if (s1 < x)
      swap(s1, x);
    s2 = max(s2, x);
  }
  void mn(ll x) {
    if (s1 == x) {
      s1 = s2;
      s2 = 0;
    }
  }
};
ll n, cs[100005], ans;
vector<ll> g[100005];
map<ll, t> dp[100005];
t dfs(ll v, ll p, ll f) {
  if (dp[v].count(p))
    return dp[v][p];
  t res;
  if (p == -1 || f) {
    for (int i = 0; i < g[v].size(); i++) {
      ll u = g[v][i];
      if (u == p)
        continue;
      t x = dfs(u, v, f);
      x.r -= 2;
      x.s1--;
      if (x.r > 0) {
        res.r += x.r;
        res.up(-x.r + x.s1);
      } else
        res.up(x.s1);
    }
    res.r += cs[v];
    res.s1 += res.r;
    res.s2 += res.r;
  } else {
    res = dfs(v, -1, 0);
    t x = dfs(p, v, 0);
    x.r -= 2;
    x.s1--;
    if (x.r > 0) {
      res.r -= x.r;
      res.mn(res.r + x.s1);
      res.s1 -= x.r;
    } else
      res.mn(res.r + x.s1);
  }
  return dp[v][p] = res;
}
int main(void) {
  scanf("%lld", &n);
  ll ma = -INF;
  for (int i = 0; i < n; i++) {
    scanf("%lld", cs + i);
    ma = max(ma, cs[i]);
  }
  if (ma <= 0) {
    printf("%lld\n", ma);
    return 0;
  }
  for (int i = 0; i < n - 1; i++) {
    ll a, b;
    scanf("%lld%lld", &a, &b);
    g[--a].PB(--b);
    g[b].PB(a);
  }
  dfs(0, -1, 1);
  for (int i = 0; i < n; i++)
    ans = max(ans, dfs(i, -1, 0).s1);
  printf("%lld\n", ans);
}