Edgy Trees

https://codeforces.com/contest/1139/problem/C
题解：并查集+快速幂
/*
*@Author:   STZG
*@Language: C++
*/
#include <bits/stdc++.h>
#include<iostream>
#include<algorithm>
#include<cstdlib>
#include<cstring>
#include<cstdio>
#include<string>
#include<vector>
#include<bitset>
#include<queue>
#include<deque>
#include<stack>
#include<cmath>
#include<list>
#include<map>
#include<set>
//#define DEBUG
#define RI register int
using namespace std;
typedef long long ll;
//typedef __int128 lll;
const int N=200000+10;
const int M=100000+10;
const int MOD=1e9+7;
const double PI = acos(-1.0);
const double EXP = 1E-8;
const int INF = 0x3f3f3f3f;
int t,n,m,k,v,u;
ll ans,cnt,flag,temp,sum[N];
int a[N];
char str;
int pre[N];
ll powl(ll a, ll n, ll p)    //快速幂 a^n % p
{
    ll ans = 1;
    while(n)
    {
        if(n & 1) ans = ans * a % p;
        a = a * a % p;
        n >>= 1;
    }
    return ans;
}
int find(int x){
    if(x==pre[x]){
        return x;
    }
    return pre[x]=find(pre[x]);
}
void add(int x,int y){
    int tx=find(x);
    int ty=find(y);
    if(tx!=ty){
        pre[tx]=ty;
        sum[ty]+=sum[tx];
    }
}
int main()
{
#ifdef DEBUG
	freopen("input.in", "r", stdin);
	//freopen("output.out", "w", stdout);
#endif
    //ios::sync_with_stdio(false);
    //cin.tie(0);
    //cout.tie(0);
    //scanf("%d",&t);
    //while(t--){
    scanf("%d%d",&n,&k);
    for(int i=1;i<=n;i++){
        pre[i]=i;
        sum[i]=1;
    }
    for(int i=1;i<n;i++){
        scanf("%d%d%d",&u,&v,&t);
        flag|=t;
        if(!t)add(u,v);
    }
    if(!flag){
        cout<<0<<endl;
        return 0;
    }
    temp=powl(n,k,MOD);//cout<<temp<<endl;
    for(int i=1;i<=n;i++){
        if(pre[i]==i){
            ans=(ans+powl(sum[i],k,MOD))%MOD;
        }
    }
    ans=(temp-ans+MOD)%MOD;
    cout<<ans<<endl;
    //}

#ifdef DEBUG
	printf("Time cost : %lf s\n",(double)clock()/CLOCKS_PER_SEC);
#endif
    //cout << "Hello world!" << endl;
    return 0;
}
猜你喜欢
