Problem:
Score : 1900 points

Problem Statement

There is a tree with N vertices.
The vertices are numbered 1 through N.
For each 1 ≤ i ≤ N - 1, the i-th edge connects vertices a_i and b_i.
The lengths of all the edges are 1.

Snuke likes some of the vertices.
The information on his favorite vertices are given to you as a string s of length N.
For each 1 ≤ i ≤ N, s_i is 1 if Snuke likes vertex i, and 0 if he does not like vertex i.

Initially, all the vertices are white.
Snuke will perform the following operation exactly once:

Select a vertex v that he likes, and a non-negative integer d. Then, paint all the vertices black whose distances from v are at most d.

Find the number of the possible combinations of colors of the vertices after the operation.

Constraints

2 ≤ N ≤ 2×10^5

1 ≤ a_i, b_i ≤ N

The given graph is a tree.

|s| = N

s consists of 0 and 1.

s contains at least one occurrence of 1.

Partial Score

In the test set worth 1300 points, s consists only of 1.

Input

The input is given from Standard Input in the following format:

N
a_1 b_1
a_2 b_2
:
a_{N - 1} b_{N - 1}
s

Output

Print the number of the possible combinations of colors of the vertices after the operation.

Sample Input 1

4
1 2
1 3
1 4
1100

Sample Output 1

4

The following four combinations of colors of the vertices are possible:

Sample Input 2

5
1 2
1 3
1 4
4 5
11111

Sample Output 2

11

This case satisfies the additional constraint for the partial score.

Sample Input 3

6
1 2
1 3
1 4
2 5
2 6
100011

Sample Output 3

8

Solution in C++14 (GCC 5.4.1) (judge verdict: Wrong Answer)

#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<llint> dp[200005], dp2[200005];
llint pid[200005];
vector<P> E;
string s;

map<P, llint> Ginv;
vector<llint> lmin[200005], rmin[200005];

void dfs(int v, int p)
{
	pid[v] = -1;
	for(int i = 0; i < G[v].size(); i++){
		int u = G[v][i];
		if(u == p){
			pid[v] = i;
			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, dp[u][j]);
			mn = min(mn, dp2[u][j]);
		}
		dp[v][i] = mx + 1;
		if(s[u] == '1') mn = 0;
		if(s[v] == '1') mn = -1;
		dp2[v][i] = 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++){
		lmx[i] = dp[v][i];
		if(i > 0) lmx[i] = max(lmx[i], lmx[i-1]);
		lmn[i] = dp2[v][i];
		if(i > 0) lmn[i] = min(lmn[i], lmn[i-1]);
	}
	for(int i = (int)G[v].size()-1; i >= 0; i--){
		rmx[i] = dp[v][i];
		if(i < (int)G[v].size()-1) rmx[i] = max(rmx[i], rmx[i+1]);
		rmn[i] = dp2[v][i];
		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]);
		dp[u][pid[u]] = 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] == '1') mn = 0;
		if(s[u] == '1') mn = -1;
		dp2[u][pid[u]] = 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;
	
	for(int i = 1; i <= n; i++) dp[i].resize(G[i].size()), dp2[i].resize(G[i].size());
	dfs(1, -1);
	dfs2(1, -1);
	
	/*for(int i = 1; i <= n; i++){
		for(int j = 0; j < dp[i].size(); j++){
			cout << i << G[i][j] << " " << dp2[i][j] << endl;
		}
	}*/
	
	llint ans = 0;
	for(int i = 1; i <= n; i++){
		//cout << i << " " << ans << endl;
		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(dp[i][j]);
			sort(vec.rbegin(), vec.rend());
			ans += min(vec[0], vec[1]);
		}
		else{
			vector<P> vec, vec2;
			if(G[i].size() <= 1) continue;
			for(int j = 0; j < G[i].size(); j++){
				llint dep = dp[i][j], dep2 = dp2[i][j], x = 1;
				if(dep >= dep2) x = 0;
				vec.push_back(make_pair(dep, x));
			}
			sort(vec.begin(), vec.end());
			
			bool flag = false;
			for(int j = 0; j < vec.size(); j++){
				if(vec[j].second == 0) flag = true;
				if(flag) vec2.push_back(vec[j]);
			}
			if(vec2.size() <= 1) continue;
			
			if(vec2.size() >= 3){
				llint mx2 = vec2[vec2.size()-2].first;
				ans += max(0LL, mx2 - vec2.front().first + 1);
			}
			/*set<llint> S;
			if(vec[1].first < mn2 && vec[1].second == 1) S.insert(vec[1].first);
			if(vec[0].first < mn2 && vec[0].second == 1) S.insert(vec[0].first);
			ans += S.size();*/
		}
		//cout << i << " " << ans << endl;
	}
	
	for(int i = 1; i <= n; i++){
		lmin[i].resize(G[i].size()), rmin[i].resize(G[i].size());
		for(int j = 0; j < G[i].size(); j++){
			Ginv[make_pair(i, G[i][j])] = j;
			lmin[i][j] = dp2[i][j];
			if(j > 0) lmin[i][j] = min(lmin[i][j], lmin[i][j-1]);
		}
		for(int j = (int)G[i].size()-1; j >= 0; j--){
			rmin[i][j] = dp2[i][j];
			if(j < (int)G[i].size()-1) rmin[i][j] = min(rmin[i][j], rmin[i][j+1]);
		}
	}
	
	for(int i = 0; i < E.size(); i++){
		llint u = E[i].first, v = E[i].second;
		llint umx = dp[v][Ginv[make_pair(v, u)]]-1, vmx = dp[u][Ginv[make_pair(u, v)]]-1;
		llint d = min(umx, vmx);
		
		llint umx2 = inf, vmx2 = inf;
		if(s[u] == '1') umx2 = 0; if(s[v] == '1') vmx2 = 0;
		llint upos = Ginv[make_pair(u, v)], vpos = Ginv[make_pair(v, u)];
		if(upos > 0) umx2 = min(umx2, lmin[u][upos-1]);
		if(upos < (int)G[u].size()-1) umx2 = min(umx2, rmin[u][upos+1]);
		if(vpos > 0) vmx2 = min(vmx2, lmin[v][vpos-1]);
		if(vpos < (int)G[v].size()-1) vmx2 = min(vmx2, rmin[v][vpos+1]);
		
		//bool flag = (umx == d && umx >= umx2) || (vmx == d && vmx >= vmx2);
		//cout << u << " " << v << " " << umx2 << " " << vmx2 << " " << flag << endl;
		
		if((umx == d && umx >= umx2) || (vmx == d && vmx >= vmx2)) ans++;
		//cout << u << " " << v << " " << flag << endl;
	}
	
	cout << ans << endl;
	return 0;
}