Submission #1046583

---

Source Code Expand

#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];
int cnt[200000];
string str;
vector<int>pat[200000];
vector<int>ko[200000];
bool flag[200000];
#define ROOT 0
typedef pair<int, int>pii;
void dfs(int node)
{
	if (flag[node])return;
	flag[node] = true;
	if (str[node] == '1')cnt[node]++;
	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]);
			cnt[node] += cnt[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;
	if (r1 >= 0)mini = min(mini, max(dp1[ko[node][r1]] - 1, m2));
	if (cnt[node] > 0)mini = min(mini, m1);
	dp1[node] = mini;
}
void calc2(int node)
{
	vector<pii>m;
	vector<int>zm, 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]]);
		zm.push_back(len1[ko[node][i]]);
	}
	if (node != ROOT)
	{
		m.push_back(make_pair(len2[node], m.size()));
		zd.push_back(dp2[node]);
		zm.push_back(len2[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;
		vector<int>z;
		for (int j = 0; j < min(3, x); j++)
		{
			if (i == m[j].second)continue;
			z.push_back(m[j].second);
		}
		if (z.size() >= 2)mini = min(mini, max(zd[z[0]] - 1, zm[z[1]]));
		if (z.size() == 1)mini = min(mini, max(zd[z[0]] - 1, 0));
		if (cnt[ROOT] - cnt[ko[node][i]] > 0 && z.size() > 0)mini = min(mini, zm[z[0]]);
		if (cnt[ROOT] - cnt[ko[node][i]] > 0 && z.size() == 0)mini = min(mini, 0);
		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;
		vector<int>zm, zd;
		for (int j = 0; j < ko[i].size(); j++)
		{
			m.push_back(make_pair(len1[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()));
			zm.push_back(len2[i]);
			zd.push_back(dp2[i]);
		}
		sort(m.begin(), m.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);
}

Submission Info

Compile Error

./Main.cpp: In function ‘int main()’:
./Main.cpp:98:19: warning: ignoring return value of ‘int scanf(const char*, ...)’, declared with attribute warn_unused_result [-Wunused-result]
  scanf("%d", &num);
                   ^
./Main.cpp:102:26: warning: ignoring return value of ‘int scanf(const char*, ...)’, declared with attribute warn_unused_result [-Wunused-result]
   scanf("%d%d", &za, &zb);
                          ^

Judge Result