Codeforces Round #548 (Div. 2), problem: (C) Edgy Trees 【并查集+快速幂】

题意

给了一棵树，n个点，m条边。让从中选k个点，使得从a1到a2，a2到a3，ak-1到ak的路径中至少经过一条黑色的边，问这样的集合有多少个

思路

用并查集统计一个连通块的节点个数，最后用总的减去他，设x是连通块的节点个数，o个联通块

code

#include<bits/stdc++.h>
#define endl '\n'
using namespace std;
typedef long long ll;
const int maxn=1e5+5;
const int mod=1e9+7;
int fa[maxn],cnt[maxn];
int n,k;
ll qpow(ll x,ll y){
    ll ba=x,ans=1;
    while(y){
        if(y&1) ans=(ans*ba)%mod;
        ba=(ba*ba)%mod;
        y>>=1;
    }
    return ans;
}
int findFa(int x){
    if(x!=fa[x])
        fa[x]=findFa(fa[x]);
    return fa[x];
}
void join(int x,int y){
    int a=findFa(x),b=findFa(y);
    if(a!=b){
        fa[a]=b;
        cnt[b]+=cnt[a];
    }
}
int main(){
    ios::sync_with_stdio(false);
    cin.tie(0);
    cin>>n>>k;
    for(int i=1;i<=n;i++)
        fa[i]=i,cnt[i]=1;
    for(int i=0;i<n-1;i++){
        int u,v,w;
        cin>>u>>v>>w;
        if(!w)
            join(u,v);
    }
    ll res=qpow(n,k);
    for(int i=1;i<=n;i++){
        if(fa[i]==i){
            res-=qpow(cnt[i],k);
            res=(res+mod)%mod;
        }
    }
    cout<<res<<endl;
    return 0;
}

学如逆水行舟，不进则退