Problem p03843 (C++)

Buggy submission (status: Runtime Error):
#include <algorithm>
#include <cctype>
#include <cstdio>
#include <cstring>
using namespace std;
typedef long long int64;
inline int read(int f = 1, int x = 0, char ch = ' ') {
  while (!isdigit(ch = getchar()))
    if (ch == '-')
      f = -1;
  while (isdigit(ch))
    x = x * 10 + ch - '0', ch = getchar();
  return f * x;
}
const int N = 1e5 + 5, INF = 0x3f3f3f3f;
struct Edge {
  int next, to;
  Edge(int next = 0, int to = 0) : next(next), to(to){};
} edge[N << 1];
int n, tot, head[N];
void _add(int x, int y) {
  edge[++tot] = Edge(head[x], y);
  head[x] = tot;
}
void add(int x, int y) {
  _add(x, y);
  _add(y, x);
}
int s[N], d1[N], d2[N], d[N];
char c[N];
int64 ans;
void dfs1(int x, int fa) {
  c[x] == '1' ? d[x] = 0, s[x] = 1 : d[x] = INF;
  for (int i = head[x]; i; i = edge[i].next) {
    int y = edge[i].to;
    if (y == fa)
      continue;
    dfs1(y, x), s[x] += s[y];
    if (d1[y] + 1 > d1[x])
      d2[x] = d1[x], d1[x] = d1[y] + 1;
    else if (d1[y] + 1 > d2[x])
      d2[x] = d1[y] + 1;
    if (s[y])
      d[x] = min(d[x], d1[y] + 1);
  }
}
void dfs2(int x, int fa) {
  int r = min(d2[x] + 1, d1[x] - 1);
  if (d[x] <= r)
    ans += r - d[x] + 1;
  for (int i = head[x]; i; i = edge[i].next) {
    int y = edge[i].to;
    if (y == fa)
      continue;
    int k = d1[y] + 1 == d1[x] ? d2[x] + 1 : d1[x] + 1;
    if (k > d1[y])
      d2[y] = d1[y], d1[y] = k;
    else if (k > d2[y])
      d2[y] = k;
    if (s[1] - s[y])
      d[y] = min(d[y], k);
    dfs2(y, x);
  }
}
int main() {
  n = read();
  for (int i = 1, x, y; i < n; ++i)
    x = read(), y = read(), add(x, y);
  scanf("%s", c + 1), dfs1(1, 0), dfs2(1, 0);
  printf("%lld\n", ans + 1);
  return 0;
}

Fixed submission:
#include <algorithm>
#include <cctype>
#include <cstdio>
#include <cstring>
using namespace std;
typedef long long int64;
inline int read(int f = 1, int x = 0, char ch = ' ') {
  while (!isdigit(ch = getchar()))
    if (ch == '-')
      f = -1;
  while (isdigit(ch))
    x = x * 10 + ch - '0', ch = getchar();
  return f * x;
}
const int N = 2e5 + 5, INF = 0x3f3f3f3f;
struct Edge {
  int next, to;
  Edge(int next = 0, int to = 0) : next(next), to(to){};
} edge[N << 1];
int n, tot, head[N];
void _add(int x, int y) {
  edge[++tot] = Edge(head[x], y);
  head[x] = tot;
}
void add(int x, int y) {
  _add(x, y);
  _add(y, x);
}
int s[N], d1[N], d2[N], d[N];
char c[N];
int64 ans;
void dfs1(int x, int fa) {
  c[x] == '1' ? d[x] = 0, s[x] = 1 : d[x] = INF;
  for (int i = head[x]; i; i = edge[i].next) {
    int y = edge[i].to;
    if (y == fa)
      continue;
    dfs1(y, x), s[x] += s[y];
    if (d1[y] + 1 > d1[x])
      d2[x] = d1[x], d1[x] = d1[y] + 1;
    else if (d1[y] + 1 > d2[x])
      d2[x] = d1[y] + 1;
    if (s[y])
      d[x] = min(d[x], d1[y] + 1);
  }
}
void dfs2(int x, int fa) {
  int r = min(d2[x] + 1, d1[x] - 1);
  if (d[x] <= r)
    ans += r - d[x] + 1;
  for (int i = head[x]; i; i = edge[i].next) {
    int y = edge[i].to;
    if (y == fa)
      continue;
    int k = d1[y] + 1 == d1[x] ? d2[x] + 1 : d1[x] + 1;
    if (k > d1[y])
      d2[y] = d1[y], d1[y] = k;
    else if (k > d2[y])
      d2[y] = k;
    if (s[1] - s[y])
      d[y] = min(d[y], k);
    dfs2(y, x);
  }
}
int main() {
  n = read();
  for (int i = 1, x, y; i < n; ++i)
    x = read(), y = read(), add(x, y);
  scanf("%s", c + 1), dfs1(1, 0), dfs2(1, 0);
  printf("%lld\n", ans + 1);
  return 0;
}