Language: C++

#include <bits/stdc++.h>
#define Inf 0x3f3f3f3f
#define rg register
#define pii pair<int, int>
#define mp(x, y) make_pair(x, y)
using namespace std;
typedef long long ll;
const int Maxn = 2e5 + 5;
typedef int arr[Maxn];
//ll Add(ll x, ll y) { return (x += y) >= Mod ? x - Mod : (x < 0 ? x + Mod : x); }
ll min(ll x, ll y) { return x < y ? x : y; }
ll max(ll x, ll y) { return x > y ? x : y; }
inline ll read()
{
	rg ll res = 0, sign = 1;
	rg char ch = getchar();
	while (!isdigit(ch))
	{
		if (ch == '-')
			sign = -1;
		ch = getchar();
	}
	do
	{
		res = (res << 3) + (res << 1) + (ch ^ 48);
		ch = getchar();
	} while (isdigit(ch));
	return res * sign;
}
int n, tot;
arr fa, deep;
int Max[Maxn][2], sum[Maxn];
ll ans;
char s[Maxn];
vector<int> to[Maxn];
void Add(int x, int p)
{
	if (deep[p] + 1 > deep[Max[x][1]])
	{
		if (deep[p] + 1 > deep[Max[x][0]])
		{
			Max[x][1] = Max[x][0];
			Max[x][0] = p;
			deep[x] = deep[Max[x][0]] + 1;
		}
		else
			Max[x][1] = p;
	}
}
void DFS(int x)
{
	for (rg int i = 0, nx; i < to[x].size(); i++)
	{
		nx = to[x][i];
		if (nx == fa[x])
			continue;
		fa[nx] = x;
		DFS(nx);
		sum[x] += sum[nx];
		Add(x, nx);
	}
}
void Solve(int x)
{
	if (to[x].size() == 1)
	{
		ans += (s[x] == '1') * 2;
		return;
	}
	ans += min(deep[x], deep[Max[x][1]] + 2);
	if (s[x] == '0')
	{
		rg int Min = Inf;
		for (rg int i = 0, nx; i < to[x].size(); i++)
		{
			nx = to[x][i];
			if (nx == fa[x] || !sum[nx])
				continue;
			if (deep[nx] < Min)
				Min = deep[nx] + 1;
		}
		ans -= Min;
	}
	rg int tmp;
	for (rg int i = 0, nx; i < to[x].size(); i++)
	{
		nx = to[x][i];
		if (nx == fa[x])
			continue;
		sum[x] = tot - sum[nx];
		tmp = deep[nx];
		deep[x] = deep[Max[x][Max[x][0] == nx]] + 1;
		Add(nx, x);
		Solve(nx);
		deep[nx] = tmp;
	}
}
int main()
{
	rg int a, b, tmp;
	cin >> n;
	for (int i = 1; i < n; i++)
	{
		a = read();
		b = read();
		to[a].push_back(b);
		to[b].push_back(a);
	}
	cin >> s + 1;
	if (n == 1)
	{
		printf("%c\n", s[1]);
		return 0;
	}
	for (int i = 1; i <= n; i++)
		tot += sum[i] = s[i] == '1';
	DFS(1);
	lim[1] = Inf;
	Solve(1);
	cout << ans + 1 << '\n';
	return 0;
}


Compiler output:
a.cc: In function 'll read()':
a.cc:15:15: warning: ISO C++17 does not allow 'register' storage class specifier [-Wregister]
   15 |         rg ll res = 0, sign = 1;
      |               ^~~
a.cc:15:24: warning: ISO C++17 does not allow 'register' storage class specifier [-Wregister]
   15 |         rg ll res = 0, sign = 1;
      |                        ^~~~
a.cc:16:17: warning: ISO C++17 does not allow 'register' storage class specifier [-Wregister]
   16 |         rg char ch = getchar();
      |                 ^~
a.cc: In function 'void DFS(int)':
a.cc:52:21: warning: ISO C++17 does not allow 'register' storage class specifier [-Wregister]
   52 |         for (rg int i = 0, nx; i < to[x].size(); i++)
      |                     ^
a.cc:52:28: warning: ISO C++17 does not allow 'register' storage class specifier [-Wregister]
   52 |         for (rg int i = 0, nx; i < to[x].size(); i++)
      |                            ^~
a.cc: In function 'void Solve(int)':
a.cc:73:24: warning: ISO C++17 does not allow 'register' storage class specifier [-Wregister]
   73 |                 rg int Min = Inf;
      |                        ^~~
a.cc:74:29: warning: ISO C++17 does not allow 'register' storage class specifier [-Wregister]
   74 |                 for (rg int i = 0, nx; i < to[x].size(); i++)
      |                             ^
a.cc:74:36: warning: ISO C++17 does not allow 'register' storage class specifier [-Wregister]
   74 |                 for (rg int i = 0, nx; i < to[x].size(); i++)
      |                                    ^~
a.cc:84:16: warning: ISO C++17 does not allow 'register' storage class specifier [-Wregister]
   84 |         rg int tmp;
      |                ^~~
a.cc:85:21: warning: ISO C++17 does not allow 'register' storage class specifier [-Wregister]
   85 |         for (rg int i = 0, nx; i < to[x].size(); i++)
      |                     ^
a.cc:85:28: warning: ISO C++17 does not allow 'register' storage class specifier [-Wregister]
   85 |         for (rg int i = 0, nx; i < to[x].size(); i++)
      |                            ^~
a.cc: In function 'int main()':
a.cc:100:16: warning: ISO C++17 does not allow 'register' storage class specifier [-Wregister]
  100 |         rg int a, b, tmp;
      |                ^
a.cc:100:19: warning: ISO C++17 does not allow 'register' storage class specifier [-Wregister]
  100 |         rg int a, b, tmp;
      |                   ^
a.cc:100:22: warning: ISO C++17 does not allow 'register' storage class specifier [-Wregister]
  100 |         rg int a, b, tmp;
      |                      ^~~
a.cc:118:9: error: 'lim' was not declared in this scope
  118 |         lim[1] = Inf;
      |         ^~~