Codeforces Round #548 (Div. 2)C. Edgy Trees 并查集

codeforces 1139C

题目链接：http://codeforces.com/contest/1139/problem/C

题意： 给你一个n个结点n-1条边的无向连通图，由红边和黑边组成。然后给一个数字k，问有多少个长度为k的结点序列（结点可重复出现），按照这些结点顺序在图上走，最少经过一次黑边。

题解

直接去算至少经过一次黑边的个数很难算，可以间接去求。

算一下所有长度为k的序列有多少个，即n的k次幂

然后算只经过红边的个数有多少，减一下就可以了。

输入边的时候判断一下是红边还是黑边，是黑边不处理，是红边那么连起来，并且在根节点处记录这个连通块有多少结点。

设某个连通块的结点有x个，则这个连通块有 x的k次幂种方案。

#include<algorithm>
#include <iostream>
#include<cstring>
#include <cstdio>
using namespace std;
typedef long long ll;
const int maxn=2e5+5;
const ll mod=1e9+7;
int father[maxn];
int num[maxn];
int get(int x){
    return x==father[x]?x:father[x]=get(father[x]);
}
ll ksm(ll a, ll b){
    a%=mod;
    ll ans=1;
    while(b){
        if(b&1) ans=ans*a%mod;
        b>>=1;
        a=a*a%mod;
    }
    return ans;
}
int main(){
    int n,k;
    cin>>n>>k;
    for(int i=1;i<=n;i++){
        father[i]=i;
        num[i]=1;
    }
    for(int i=1;i<n;i++){
        int a,b,x;
        scanf("%d%d%d",&a,&b,&x);
        if(x==0){
            int c=get(a);
            int d=get(b);
            if(c!=d){
                father[d]=c;
                num[c]+=num[d];
                num[d]=0;
            }
        }
    }
    ll ans=ksm(n,k);
    for(int i=1;i<=n;i++){
        if(num[i]){
            ans=(ans-ksm(num[i],k)+mod)%mod;
        }
    }
    cout<<ans<<endl;
    return 0;
}