Problem p03843 (C++)

Buggy submission (status: Runtime Error):
#include <bits/stdc++.h>
using namespace std;
#define pb push_back
#define SZ(x) (int((x).size()))
#define rep(i, a, b) for (int(i) = (a); (i) <= (b); (i)++)
#define rep2(i, a, b) for (int(i) = (a); (i) < (b); (i)++)
#define Rep(p, x) for (int(p) = head[(x)]; (p); (p) = nxt[(p)])
const int maxn = 1e5 + 5;
long long ans;
int n, tot, all;
int f[maxn], g[maxn], cnt[maxn], mx[maxn], se[maxn], high[maxn], low[maxn];
char s[maxn];
vector<int> e[maxn];
void dfs1(int x, int fa) {
  mx[x] = f[x] = 0;
  cnt[x] = (s[x] == '1');
  rep2(i, 0, SZ(e[x])) {
    int y = e[x][i];
    if (y == fa)
      continue;
    dfs1(y, x);
    f[x] = max(f[y] + 1, f[x]);
    if (f[y] + 1 > mx[x])
      se[x] = mx[x], mx[x] = f[y] + 1;
    else if (f[y] + 1 > se[x])
      se[x] = f[y] + 1;
    cnt[x] += cnt[y];
  }
}
void dfs2(int x, int fa) {
  rep2(i, 0, SZ(e[x])) {
    int y = e[x][i];
    if (y == fa)
      continue;
    g[y] = max((f[y] + 1 == mx[x] ? se[x] : mx[x]), g[x]) + 1;
    dfs2(y, x);
  }
}
void dfs3(int x, int fa) {
  if (g[x] > mx[x])
    se[x] = mx[x], mx[x] = g[x];
  else if (g[x] > se[x])
    se[x] = g[x];
  high[x] = min(se[x] + 1, mx[x] - 1);
  low[x] = high[x] + 1;
  rep2(i, 0, SZ(e[x])) {
    int y = e[x][i];
    if (y == fa)
      continue;
    dfs3(y, x);
    if (cnt[y])
      low[x] = min(low[x], f[y] + 1);
  }
  if (all - cnt[x])
    low[x] = min(low[x], g[x]);
  if (s[x] == '1')
    low[x] = 0;
}
int main() {
  scanf("%d", &n);
  rep(i, 1, n - 1) {
    int x, y;
    scanf("%d%d", &x, &y);
    e[x].pb(y);
    e[y].pb(x);
  }
  scanf("%s", s + 1);
  rep(i, 1, n) all += (s[i] == '1');
  dfs1(1, 0);
  dfs2(1, 0);
  dfs3(1, 0);
  rep(x, 1, n) ans += high[x] - low[x] + 1;
  printf("%lld\n", ans + 1);
  return 0;
}

Fixed submission:
#include <bits/stdc++.h>
using namespace std;
#define pb push_back
#define SZ(x) (int((x).size()))
#define rep(i, a, b) for (int(i) = (a); (i) <= (b); (i)++)
#define rep2(i, a, b) for (int(i) = (a); (i) < (b); (i)++)
#define Rep(p, x) for (int(p) = head[(x)]; (p); (p) = nxt[(p)])
const int maxn = 2e5 + 5;
long long ans;
int n, tot, all;
int f[maxn], g[maxn], cnt[maxn], mx[maxn], se[maxn], high[maxn], low[maxn];
char s[maxn];
vector<int> e[maxn];
void dfs1(int x, int fa) {
  mx[x] = f[x] = 0;
  cnt[x] = (s[x] == '1');
  rep2(i, 0, SZ(e[x])) {
    int y = e[x][i];
    if (y == fa)
      continue;
    dfs1(y, x);
    f[x] = max(f[y] + 1, f[x]);
    if (f[y] + 1 > mx[x])
      se[x] = mx[x], mx[x] = f[y] + 1;
    else if (f[y] + 1 > se[x])
      se[x] = f[y] + 1;
    cnt[x] += cnt[y];
  }
}
void dfs2(int x, int fa) {
  rep2(i, 0, SZ(e[x])) {
    int y = e[x][i];
    if (y == fa)
      continue;
    g[y] = max((f[y] + 1 == mx[x] ? se[x] : mx[x]), g[x]) + 1;
    dfs2(y, x);
  }
}
void dfs3(int x, int fa) {
  if (g[x] > mx[x])
    se[x] = mx[x], mx[x] = g[x];
  else if (g[x] > se[x])
    se[x] = g[x];
  high[x] = min(se[x] + 1, mx[x] - 1);
  low[x] = high[x] + 1;
  rep2(i, 0, SZ(e[x])) {
    int y = e[x][i];
    if (y == fa)
      continue;
    dfs3(y, x);
    if (cnt[y])
      low[x] = min(low[x], f[y] + 1);
  }
  if (all - cnt[x])
    low[x] = min(low[x], g[x]);
  if (s[x] == '1')
    low[x] = 0;
}
int main() {
  scanf("%d", &n);
  rep(i, 1, n - 1) {
    int x, y;
    scanf("%d%d", &x, &y);
    e[x].pb(y);
    e[y].pb(x);
  }
  scanf("%s", s + 1);
  rep(i, 1, n) all += (s[i] == '1');
  dfs1(1, 0);
  dfs2(1, 0);
  dfs3(1, 0);
  rep(x, 1, n) ans += high[x] - low[x] + 1;
  printf("%lld\n", ans + 1);
  return 0;
}