Submission #2119095

---

Source Code Expand

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef pair<int,int> P;
typedef pair<int,P> P1;
typedef pair<P,P> P2;
#define pb push_back
#define mp make_pair
#define eps 1e-7
#define INF 1000000000
#define mod 1000000007
#define fi first
#define sc second
#define rep(i,x) for(int i=0;i<x;i++)
#define repn(i,x) for(int i=1;i<=x;i++)
#define SORT(x) sort(x.begin(),x.end())
#define ERASE(x) x.erase(unique(x.begin(),x.end()),x.end())
#define POSL(x,v) (lower_bound(x.begin(),x.end(),v)-x.begin())
#define POSU(x,v) (upper_bound(x.begin(),x.end(),v)-x.begin())
map<P,int>dp;
map<P,bool>ex;
vector<int>edge[200005];
int n;
bool ok[200005];
void dfs(int v,int u){
	if(u!=-1){
		dp[mp(u,v)] = 1;
		ex[mp(u,v)] = ok[v];
	}
	rep(i,edge[v].size()){
		if(edge[v][i]==u) continue;
		dfs(edge[v][i],v);
		if(u!=-1){
			dp[mp(u,v)] = max(dp[mp(u,v)],dp[mp(v,edge[v][i])]+1);
			ex[mp(u,v)] |= ex[mp(v,edge[v][i])];
		}
	}
}
void dfs2(int v,int u,int mx,bool EX){
	EX |= ok[v];
	priority_queue<int,vector<int>,greater<int> >S;
	int cnt = 0;
	rep(i,edge[v].size()){
		if(edge[v][i]==u) continue;
		S.push(dp[mp(v,edge[v][i])]);
		if(S.size() == 3) S.pop();
		cnt += ex[mp(v,edge[v][i])];
	}
	if(S.size() < 1) return;
	S.push(0);
	if(S.size() == 3) S.pop();
	int a = S.top(); S.pop();
	int b = S.top(); S.pop();
	rep(i,edge[v].size()){
		if(edge[v][i]==u) continue;
		if(dp[mp(v,edge[v][i])] == b) dp[mp(edge[v][i],v)] = 1+max(mx,a);
		else dp[mp(edge[v][i],v)] = 1+max(mx,b);
		if(EX || cnt-ex[mp(v,edge[v][i])] >= 1){
			ex[mp(edge[v][i],v)] = 1;
		}
		else{
			ex[mp(edge[v][i],v)] = 0;
		}
		dfs2(edge[v][i],v,dp[mp(edge[v][i],v)],ex[mp(edge[v][i],v)]);
	}
}
int main(){
	scanf("%d",&n);
	rep(i,n-1){
		int a,b; scanf("%d%d",&a,&b);
		a--; b--;
		edge[a].pb(b);
		edge[b].pb(a);
	}
	string s; cin>>s;
	rep(i,n) if(s[i]=='1') ok[i]=true;
	dfs(0,-1);
	dfs2(0,-1,0,false);
	ll ret = 0;
	rep(i,n){
		priority_queue<int,vector<int>,greater<int> >S;
		for(int j=0;j<edge[i].size();j++){
			S.push(dp[mp(i,edge[i][j])]);
			if(S.size() == 3) S.pop();
		}
		int ub = 0;
		if(S.size() == 2){
			ub = max(ub,S.top());
		}
		int lb = INF;
		if(ok[i]) lb = 0;
		for(int j=0;j<edge[i].size();j++){
			if(!ex[mp(i,edge[i][j])]) continue;
			lb = min(lb,dp[mp(i,edge[i][j])]);
		}
		if(lb <= ub){
			ret += (ub-lb+1);
		}
	}
	rep(i,n){
		rep(j,edge[i].size()){
			int u = i,v = edge[i][j];
			if(u > v) continue;
			bool flag = 0;
			if(dp[mp(v,u)] <= dp[mp(u,v)]){
				flag |= ex[mp(v,u)];
			}
			if(dp[mp(v,u)] >= dp[mp(u,v)]){
				flag |= ex[mp(u,v)];
			}
			if(flag) ret++;
		}
	}
	cout << ret << endl;
}

Submission Info

Compile Error

./Main.cpp: In function ‘int main()’:
./Main.cpp:68:16: warning: ignoring return value of ‘int scanf(const char*, ...)’, declared with attribute warn_unused_result [-Wunused-result]
  scanf("%d",&n);
                ^
./Main.cpp:70:31: warning: ignoring return value of ‘int scanf(const char*, ...)’, declared with attribute warn_unused_result [-Wunused-result]
   int a,b; scanf("%d%d",&a,&b);
                               ^

Judge Result