Language: C++

#include<stdio.h>
#include<vector>
#include<algorithm>
#include<string>
#include<iostream>
using namespace std;
typedef long long ll;
int dp1[200000], dp2[200000];
int len1[200000], len2[200000];
string str;
vector<int>pat[200000];
vector<int>ko[200000];
bool flag[200000];
#define ROOT (num-1)
typedef pair<int, int>pii;
void dfs(int node)
{
	if (flag[node])return;
	flag[node] = true;
	for (int i = 0; i < pat[node].size(); i++)
	{
		if (!flag[pat[node][i]])
		{
			ko[node].push_back(pat[node][i]);
			dfs(pat[node][i]);
		}
	}
}
void calc1(int node)
{
	for (int i = 0; i < ko[node].size(); i++)calc1(ko[node][i]);
	int m1 = 0, m2 = 0, r1 = -1, r2 = -2;
	for (int i = 0; i < ko[node].size(); i++)
	{
		int v = ko[node][i];
		if (m1 < len1[v])
		{
			m2 = m1, r2 = r1;
			m1 = len1[v], r1 = i;
		}
		else if (m2 < len1[v])m2 = len1[v], r2 = i;
	}
	len1[node] = m1 + 1;
	int mini = 1000000000;
	for (int i = 0; i < ko[node].size(); i++)
	{
		if (r1 != i)mini = min(mini, max(dp1[ko[node][i]] - 1, m1));
		else mini = min(mini, max(dp1[ko[node][i]] - 1, m2));
	}
	if (str[node] == '1')mini = min(mini, m1);
	dp1[node] = mini;
}
void calc2(int node)
{
	vector<pii>m;
	vector<int>zd;
	for (int i = 0; i < ko[node].size(); i++)
	{
		m.push_back(make_pair(len1[ko[node][i]], i));
		zd.push_back(dp1[ko[node][i]]);
	}
	if (node != ROOT)
	{
		m.push_back(make_pair(len2[node], m.size()));
		zd.push_back(dp2[node]);
	}
	sort(m.begin(), m.end());
	reverse(m.begin(), m.end());
	int x = m.size();
	for (int i = 0; i < ko[node].size(); i++)
	{
		int mini = 1000000000;
		for (int j = 0; j < min(3, x); j++)
		{
			if (i == m[j].second)continue;
			int maxi = 0;
			for (int k = 0; k < min(3, x); k++)
			{
				if (i == m[k].second || m[j].second == m[k].second)continue;
				maxi = max(maxi, m[k].first);
			}
			mini = min(mini, max(zd[m[j].second] - 1, maxi));
		}
		if (str[node] == '1')
		{
			if (m.size() == 1)mini = 0;
			else if (m[0].second != i)mini = min(mini, m[0].first);
			else mini = min(mini, m[1].first);
		}
		dp2[ko[node][i]] = mini;
		int maxi = 0;
		for (int j = 0; j < min(3, x); j++)
		{
			if (i == m[j].second)continue;
			maxi = max(maxi, m[j].first);
		}
		len2[ko[node][i]] = maxi + 1;
	}
	for (int i = 0; i < ko[node].size(); i++)calc2(ko[node][i]);
}
int main()
{
	int num;
	scanf("%d", &num);
	for (int i = 0; i < num - 1; i++)
	{
		int za, zb;
		scanf("%d%d", &za, &zb);
		za--, zb--;
		pat[za].push_back(zb);
		pat[zb].push_back(za);
	}
	cin >> str;
	dfs(ROOT);
	calc1(ROOT);
	calc2(ROOT);
	ll ans = 1;
	//for (int i = 0; i < num; i++)printf("%d %d   %d %d\n", len1[i], len2[i], dp1[i], dp2[i]);
	for (int i = 0; i < num; i++)
	{
		vector<pii>m, d;
		vector<int>zm, zd;
		for (int j = 0; j < ko[i].size(); j++)
		{
			m.push_back(make_pair(len1[ko[i][j]],j));
			d.push_back(make_pair(dp1[ko[i][j]],j));
			zm.push_back(len1[ko[i][j]]);
			zd.push_back(dp1[ko[i][j]]);
		}
		if (i != ROOT)
		{
			m.push_back(make_pair(len2[i], m.size()));
			d.push_back(make_pair(dp2[i], d.size()));
			zm.push_back(len2[i]);
			zd.push_back(dp2[i]);
		}
		sort(m.begin(), m.end());
		sort(d.begin(), d.end());
		reverse(m.begin(), m.end());
		ll ret = 0;
		if (m.size() >= 2)if (str[i] == '1')ret += m[1].first;
		int mini = 1000000000;
		if (m.size() >= 3)
		{
			for (int j = 0; j < m.size(); j++)
			{
				vector<int>z;
				for (int k = 0; k < 3; k++)
				{
					if (m[k].second != j)z.push_back(m[k].first);
				}
				int lb = max(zd[j] - 1, 0), ub = min(z[1] - 1, zm[j] - 1);
				ret += max(0, ub - lb + 1);
				mini = min(mini, max(lb, zm[j]));
			}
			if (str[i] == '0')ret += max(0, m[1].first - mini);
		}
		ll sum = 0;
		for (int j = 0; j < m.size(); j++)if (zd[j] == 0)sum++;
		ans += ret + sum;
		//printf("%d %lld %lld\n", i + 1, ret, sum);
	}
	printf("%lld\n", ans);
}

Compiler output:
a.cc: In function 'void calc2(int)':
a.cc:14:15: error: 'num' was not declared in this scope; did you mean 'enum'?
   14 | #define ROOT (num-1)
      |               ^~~
a.cc:62:21: note: in expansion of macro 'ROOT'
   62 |         if (node != ROOT)
      |                     ^~~~