HDU 6832 A Very Easy Graph Problem

A Very Easy Graph Problem

题意：
给了一张连通图，有黑白两种颜色的点，求所有黑白点对之间的距离之和。
对于给的边的顺序，第i条边的权值为2^i

思路：
不难知道2 ^ 1 + 2 ^ 2 +2 ^ 3 + … + 2 ^ (i-1) < 2 ^ i
所以最后的图一定是原图的最小生成树，用并查集保存下来选择的边即可。
然后我们计算在最小生成树中的边对答案的贡献。
对于边(u,v)，我们只需要知道从1到u有多少个黑点，从v到n有多少个白点，贡献就是左边黑点个数乘右边白点个数乘上边的权值，那么同理要处理左边白点，右边黑点的贡献即可。
那么黑白点的个数怎么求呢？我们只需要对这棵树进行一次dfs
cnt[x][0]表示以x为结点的子树的白点个数有多少
cnt[x][1]表示以x为结点的子树的黑点个数有多少
最后对生成树中的边，进行计算即可

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const ll mod=1e9+7;
const int N=3e5+5;
ll q[N],f[N];
int find(int x){
    return f[x]==x?x:f[x]=find(f[x]);
}
struct node{
    int x,y;
}edge[N];
struct Edge{
    int to,next;
}e[N];
int h[N],tot;
void add(int x,int y){
    e[tot]={y,h[x]},h[x]=tot++;
}
bool vis[N],id[N];
ll cnt[N][2];
int dfn[N],num;
void dfs(int x,int f){
    cnt[x][1]=id[x]^0;
    cnt[x][0]=id[x]^1;
    dfn[x]=++num;
    for(int i=h[x];~i;i=e[i].next){
        if(e[i].to==f) continue;
        dfs(e[i].to,x);
        cnt[x][1]+=cnt[e[i].to][1];
        cnt[x][0]+=cnt[e[i].to][0];
    }
}
int main(){
    q[0]=1;
    for(int i=1;i<N;i++) q[i]=q[i-1]*2%mod;
    ios::sync_with_stdio(0);
    int T;cin>>T;
    while(T--){
        int n,m;cin>>n>>m;
        tot=0,num=0;
        for(int i=1;i<=n;i++) f[i]=i,h[i]=-1;///并查集初始化 前向星初始化
        for(int i=1;i<=m;i++) vis[i]=0;///记录边是否在生成树中
        for(int i=1;i<=n;i++)  cin>>id[i];///记录点的黑白颜色
        
        for(int i=1;i<=m;i++){
            int x,y;cin>>x>>y;
            edge[i]={x,y};
            int fx=find(x),fy=find(y);
            if(fx==fy) continue;
            f[fx]=fy;
            vis[i]=1;
            add(x,y),add(y,x);
        }
        dfs(1,0);
        ll ans=0;
        for(int i=1;i<=m;i++){
            if(!vis[i]) continue;
            int fa=edge[i].x,son=edge[i].y;
            if(dfn[fa]>dfn[son]) swap(fa,son);///注意父子关系 因为是无向边
            ans+=cnt[son][1]*(cnt[1][0]-cnt[son][0])%mod*q[i]%mod;
            ans%=mod;
            ans+=cnt[son][0]*(cnt[1][1]-cnt[son][1])%mod*q[i]%mod;
            ans%=mod;
        }
        cout<<ans<<endl;
    }
    return 0;
}