2020杭电暑假多校第六场 A Very Easy Graph Problem(换根dp)

题目链接
题意给一个无向图，每个点为0或1，第i条边长度为2的i次方。题目要求出所有点i、j，i为0且j为1，i、j间最短路的和。
思路首先容易看出，若第i条边加入后图中产生了环，由于二进制的特性第i条边必然不能在最短路中，所以转化为树上问题。先计算两个值，1.每个节点到标记为1的子节点的路径长度和。2.每个节点标记为1的子节点个数。从根dfs下去换根，每次维护每个节点到所有标记为1节点的长度和。
代码

#include<bits/stdc++.h>
#define ll long long
#define LL long long
#define PB push_back
#define MP make_pair
using namespace std;
const int maxn=2e5+100;
const ll inf=1e18+10;
ll mod=1e9+7;
int n,m,f[maxn],a[maxn],num1;
ll len[maxn],num[maxn],ans=0,two[maxn];
struct node{
	ll u,to,length;
};
vector<node>g[maxn];
int F(int x){
	if(x==f[x])return x;
	return f[x]=F(f[x]);
}

inline void dfs1(int pre,int u){
	int sz=g[u].size();
	for(int i=0;i<sz;i++){
		ll to=g[u][i].to,v=g[u][i].length;
		if(to==pre)continue;
		dfs1(u,to);
		ll down_num=a[to]+num[to];
		num[u]+=down_num;
		len[u]=(len[u]+(len[to]+(v*down_num)%mod)%mod)%mod;
	}
	return ;
}

inline void dfs2(int pre,int u){
	if(!a[u])ans=(ans+len[u])%mod;
	int sz=g[u].size();
	for(int i=0;i<sz;i++){
		ll to=g[u][i].to,v=g[u][i].length;
		if(to==pre)continue;
		ll down_num=num[to]+a[to];//u下方1个数
		ll up_len=((len[u]-len[to]+mod)%mod-(v*down_num)%mod+mod)%mod;//父节点减去子节点所有 
		ll up_num=num1-down_num;//to上方1
		len[to]=(len[to]+(up_len+(v*up_num)%mod)%mod)%mod;
		dfs2(u,to);
	}
	return ;
}
void Slove(){
	scanf("%d%d",&n,&m);
	ans=num1=0;
	for(int i=1;i<=n;++i){
		f[i]=i;g[i].clear();len[i]=num[i]=0;
	}
	for(int i=1;i<=n;++i){
		scanf("%d",&a[i]);
		num1+=a[i];
	}
	for(int i=1;i<=m;i++){
		int x,y;
		scanf("%d%d",&x,&y);
		int t1=F(x),t2=F(y);
		if(t1!=t2){
			f[t1]=t2;
			g[x].PB(node{x,y,two[i]});
			g[y].PB(node{y,x,two[i]});
		}
	}
	dfs1(0,1);
	dfs2(0,1);
	printf("%lld\n",ans);
}
int main(){
	ll length=1;
	for(int i=1;i<=maxn-5;i++){
		length=length*2%mod;
		two[i]=length;
	}
	int T;
	scanf("%d",&T);
	while(T--){
		Slove();
	}
	return 0; 
}