wa on #18 只有一个点

#include<bits/stdc++.h>
using namespace std;
/*找出以谁为根*/
vector<int>u[200010];
int f[200010];
int d[200010];
int v[200010];
int y[200010];
int n,k;
int g;
int sum;
int ans;
int dfs1(int id,int fa){
	if(g){
		return 0;
	}
	int sum=0;
	for(auto i:u[id]){
		if(i!=fa){
			sum+=dfs1(i,id);
		}
	}
	d[id]=1;
	if(g){
		return 0;
	}
	++sum;
	if(sum>=k){
		int flag=0;
		for(int i=1;i<=k;i++){
			if(!d[v[i]]){
				flag=1;
				break;
			}
		}
		if(!flag){
			g=id;
		}
		//cout<<g<<endl;
    }
	return sum;
}
int dfs2(int id,int fa,int cnt){
	int sum=0;
	//cout<<id<<endl;
	if(f[id]){
		ans+=cnt;
		--cnt;
	}
	for(auto i:u[id]){
		if(i!=fa){
			sum+=dfs2(i,id,cnt+1);
		}
	}
	if(sum>=2 && !f[id]){
		ans-=sum-1;
	}
	if(f[id]){
		sum++;
	}
	return sum;
}
int main(){ 
	cin>>n;
	for(int i=1;i<=n;i++){
		cin>>y[i];
		if(y[i]){
			v[++k]=i;
			f[v[k]]=1;
		}
	}
	for(int i=1;i<n;i++){
		int a,b;
		cin>>a>>b;
		u[a].push_back(b);
		u[b].push_back(a);
	}
	dfs1(1,0);
	dfs2(g,0,1);
	cout<<ans-k;
	return 0;
}