Language: C++

#include <iostream>
#include <vector>
#include <utility>
#include <map>
#include <set>
#include <string>
#include <algorithm>
#define llint long long
#define inf 1000000000

using namespace std;
typedef pair<llint, llint> P;

llint n;
vector<llint> G[200005];
vector<P> E;
string s;
map<P, llint> mp, mp2;

void dfs(int v, int p)
{
	for(int i = 0; i < G[v].size(); i++){
		int u = G[v][i];
		if(u == p) continue;
		dfs(u, v);
		
		llint mx = 0, mn = inf;
		for(int j = 0; j < G[u].size(); j++){
			if(G[u][j] == v) continue;
			mx = max(mx, mp[make_pair(u, G[u][j])]);
			mn = min(mn, mp2[make_pair(u, G[u][j])]);
		}
		mp[make_pair(v, u)] = mx + 1;
		if(s[u] == '1') mn = 0;
		if(s[v] == '1') mn = -1;
		mp2[make_pair(v, u)] = mn + 1;
	}
}

llint lmx[200005], rmx[200005];
llint lmn[200005], rmn[200005];
void dfs2(int v, int p)
{
	for(int i = 0; i < G[v].size(); i++){
		int u = G[v][i];
		lmx[i] = mp[make_pair(v, u)];
		if(i > 0) lmx[i] = max(lmx[i], lmx[i-1]);
		lmn[i] = mp2[make_pair(v, u)];
		if(i > 0) lmn[i] = min(lmn[i], lmn[i-1]);
	}
	for(int i = (int)G[v].size()-1; i >= 0; i--){
		int u = G[v][i];
		rmx[i] = mp[make_pair(v, u)];
		if(i < (int)G[v].size()-1) rmx[i] = max(rmx[i], rmx[i+1]);
		rmn[i] = mp2[make_pair(v, u)];
		if(i < (int)G[v].size()-1) rmn[i] = min(rmn[i], rmn[i+1]);
	}
	for(int i = 0; i < G[v].size(); i++){
		int u = G[v][i];
		if(u == p) continue;
		llint mx = 0;
		if(i > 0) mx = max(mx, lmx[i-1]);
		if(i < (int)G[v].size()-1) mx = max(mx, rmx[i+1]);
		mp[make_pair(u, v)] = mx + 1;
		
		llint mn = inf;
		if(i > 0) mn = min(mn, lmn[i-1]);
		if(i < (int)G[v].size()-1) mn = min(mn, rmn[i+1]);
		if(s[v] == '0') mn = 0;
		if(s[u] == '1') mn = -1;
		mp2[make_pair(u, v)] = mn + 1;
	}
	for(int i = 0; i < G[v].size(); i++){
		if(G[v][i] == p) continue;
		dfs2(G[v][i], v);
	}
}

int main(void)
{
	ios::sync_with_stdio(0);
	cin.tie(0);
	
	cin >> n;
	llint u, v;
	for(int i = 1; i <= n-1; i++){
		cin >> u >> v;
		G[u].push_back(v);
		G[v].push_back(u);
		E.push_back(make_pair(u, v));
	}
	cin >> s;
	s = "#" + s;
	
	dfs(1, -1);
	dfs2(1, -1);
	
	/*for(auto it = mp2.begin(); it != mp2.end(); it++){
		cout << it->first.first << " " << it->first.second << " " << it->second << endl;
	}*/
	
	llint ans = 0;
	for(int i = 1; i <= n; i++){
		if(s[i] == '1'){
			ans++;
			vector<llint> vec;
			if(G[i].size() <= 1) continue;
			for(int j = 0; j < G[i].size(); j++) vec.push_back(mp[make_pair(i, G[i][j])]);
			sort(vec.rbegin(), vec.rend());
			ans += min(vec[0], vec[1]);
		}
		else{
			vector<llint> vec; set<llint> S;
			if(G[i].size() <= 1) continue;
			for(int j = 0; j < G[i].size(); j++){
				llint dep = mp[make_pair(i, G[i][j])], dep2 = mp2[make_pair(i, G[i][j])];
				if(dep <= dep2) S.insert(dep);
				vec.push_back(mp[make_pair(i, G[i][j])]);
			}
			sort(vec.rbegin(), vec.rend());
			
			bool flag = (vec[0] == vec[1]);
			vec.erase(unique(vec.begin(), vec.end()), vec.end());
			llint mx = vec[0];
			for(int j = 0; j < vec.size(); j++){
				if(vec[i].first == mx){
					if(flag) ans++;
				}
				else if(vec[i].second == 1) ans++;
			}
		}
	}
	for(int i = 0; i < E.size(); i++){
		llint u = E[i].first, v = E[i].second;
		llint umx = 0, vmx = 0;
		for(int j = 0; j < G[u].size(); j++){
			if(G[u][j] == v) continue;
			umx = max(umx, mp[make_pair(u, G[u][j])]);
		}
		for(int j = 0; j < G[v].size(); j++){
			if(G[v][j] == u) continue;
			vmx = max(vmx, mp[make_pair(v, G[v][j])]);
		}
		llint d = min(umx, vmx);
		//cout << d << endl;
		bool flag = false;
		for(int j = 0; j < G[u].size(); j++){
			if(G[u][j] == v) continue;
			if(mp[make_pair(u, G[u][j])] <= d && mp2[make_pair(u, G[u][j])] <= d) flag = true;
		}
		for(int j = 0; j < G[v].size(); j++){
			if(G[v][j] == u) continue;
			if(mp[make_pair(v, G[v][j])] <= d && mp2[make_pair(v, G[v][j])] <= d) flag = true;
		}
		if(umx == 0 && s[u] == '1' || vmx == 0 && s[v] == '1') flag = true;
		if(flag) ans++;
		//cout << u << " " << v << " " << flag << endl;
	}
	
	cout << ans << endl;
	return 0;
}

Compiler output:
a.cc: In function 'int main()':
a.cc:126:43: error: request for member 'first' in 'vec.std::vector<long long int>::operator[](((std::vector<long long int>::size_type)i))', which is of non-class type '__gnu_cxx::__alloc_traits<std::allocator<long long int>, long long int>::value_type' {aka 'long long int'}
  126 |                                 if(vec[i].first == mx){
      |                                           ^~~~~
a.cc:129:48: error: request for member 'second' in 'vec.std::vector<long long int>::operator[](((std::vector<long long int>::size_type)i))', which is of non-class type '__gnu_cxx::__alloc_traits<std::allocator<long long int>, long long int>::value_type' {aka 'long long int'}
  129 |                                 else if(vec[i].second == 1) ans++;
      |                                                ^~~~~~