Submission #1694903

---

Source Code Expand

#include <iostream>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <algorithm>
#include <vector>
#include <queue>
#include <stack>
#include <cmath>

using namespace std;

const int MAXN = 200000 + 1000;

int n, m;
long long ans = 0;
vector<int> edg[MAXN];
bool iss[MAXN];
vector<int> spe;
int fa[MAXN], dup[MAXN], ddo[MAXN], cnt[MAXN];

void dp1(int x, int father)
{
	ddo[x] = 1, cnt[x] = iss[x], fa[x] = father;
	for(int i = 0; i < edg[x].size(); i++)
		if(edg[x][i] != father)
		{
			int y = edg[x][i];
			dp1(y, x);
			ddo[x] = max(ddo[x], ddo[y] + 1);
			cnt[x] += cnt[y];
		}
}

typedef pair<int, int> pii;

pii max(const pii &a, int b)
{
	if(b < a.second) return a;
	else if(b > a.first) return make_pair(b, a.first);
	else return make_pair(a.first, b);
}

void dp2(int x, int father)
{
	pii mx(0, 0);
	for(int i = 0; i < edg[x].size(); i++)
		if(edg[x][i] != father)
			mx = max(mx, ddo[edg[x][i]]);
	for(int i = 0; i < edg[x].size(); i++)
		if(edg[x][i] != father)
		{
			int y = edg[x][i];
			dup[y] = max(dup[x], mx.first == ddo[y] ? mx.second : mx.first) + 1;
			dp2(y, x);
		}
}

void work()
{
	dp1(1, 0);
	dp2(1, 0);
	// for(int i = 1; i <= n; i++)
	// 	cout << i << ' ' << ddo[i] << ' ' << dup[i] << ' ' << cnt[i] << endl;
	for(int i = 1; i <= n; i++)
	{
		pii mx(0, 0);
		for(int j = 0; j < edg[i].size(); j++)
			if(edg[i][j] != fa[i])
				mx = max(mx, ddo[edg[i][j]]);
		mx = max(mx, dup[i]);
		int high = min(mx.second + 1, mx.first - 1);
		int low = high + 1;
		if(!iss[i])
		{
			for(int j = 0; j < edg[i].size(); j++)
				if(edg[i][j] != fa[i])
					if(cnt[edg[i][j]])
						low = min(low, ddo[edg[i][j]]);
			if(m - cnt[i])
				low = min(low, dup[i]);
		}
		else
			low = 0;
		ans += high - low + 1;
	}
	ans = ans + 1;
	cout << ans << endl;
}

int main()
{
	ios::sync_with_stdio(false);
	// freopen("1.in", "r", stdin);
	// freopen("1.out", "w", stdout);

	cin >> n;
	for(int i = 1; i <= n - 1; i++)
	{
		int x, y;
		cin >> x >> y;
		edg[x].push_back(y), edg[y].push_back(x);
	}
	static char ss[MAXN];
	cin >> (ss + 1);
	for(int i = 1; i <= n; i++)
		if(ss[i] == '1')
			iss[i] = true, spe.push_back(i);
	m = spe.size();

	work();


	return 0;
}

Submission Info

Judge Result