CF1156D 0-1-Tree 换根dp or 并查集

这题好像做法挺多的  还有个并查集的方法也非常好理解 可以看博客：CF1156D 0-1-Tree

我的做法是换根dp   由题目可知 路径是有向的  也就是说 如果 u->v满足条件 且v->u满足条件  那么我们都要算进去   那我们先算以1为出发点  可以得到的答案数 然后再进行换根  把这些答案都加起来就是最终答案了   

考虑以u为出发点  dp[u][0] 表示与子节点连边为0的方案数  dp[u][1]同理

那么状态转移就是

     

我们用这个式子算出以1为出发点的答案  再进行换根dp就可以了

#include<bits/stdc++.h>
using namespace std;
const int N = 2e5+10;
typedef long long ll;
int h[N],nex[N<<1],to[N<<1],edge[N<<1],cur;
ll dp[N][2],ans;
inline int in(){
	int x=0;char c=0;
	while(c<'0'||c>'9') c=getchar();
	while(c<='9'&&c>='0') x=(x<<1)+(x<<3)+(c^48),c=getchar();
	return x;
}
void add_edge(int x,int y,int w){
	to[++cur]=y;nex[cur]=h[x];edge[cur]=w;h[x]=cur; 
}
void dfs1(int u,int f){
	
	for(int i = h[u]; i; i = nex[i]){
		int v = to[i];
		if(v!=f){
			dfs1(v,u);
			if(edge[i]) 
			dp[u][1]+=dp[v][0]+dp[v][1]+1;
			else 
			dp[u][0]+=dp[v][0]+1;
		}
	}
	//printf("u=%d dp[u][0]=%lld dp[u][1]=%lld\n",u,dp[u][0],dp[u][1]);
}
void dfs2(int u,int f){
	//printf("u=%d dp[u][0]=%lld dp[u][1]=%lld\n",u,dp[u][0],dp[u][1]);
	ans+=dp[u][0]+dp[u][1];
	for(int i = h[u]; i; i = nex[i]){
		int v = to[i];
		if(v!=f){
			ll x=dp[v][1],y=dp[v][0];
			if(edge[i]){
				dp[u][1]-=x+y+1;
				dp[v][1]+=dp[u][1]+dp[u][0]+1;
			}else{
				dp[u][0]-=y+1;
				dp[v][0]+=dp[u][0]+1;
			}
			dfs2(v,u);
			if(edge[i]){
				dp[v][1]-=dp[u][1]+dp[u][0]+1;
				dp[u][1]+=x+y+1;
			}else{
				dp[v][0]-=dp[u][0]+1;
				dp[u][0]+=y+1;
			}
		}
	}
}
int main(){
	int n;
	n=in();
	for(int i = 1; i <= n-1; i++){
		int u,v,w;
		u=in(),v=in(),w=in();
		add_edge(u,v,w);add_edge(v,u,w);
	}
	dfs1(1,0);dfs2(1,0);
	printf("%lld\n",ans);
	return 0;
}