Codeforces-Round#548(Div.2)-C-Edgy Trees-快速冪

這個題沒想出來，好菜QAQ...

題目要求至少經過一段黑色的邊的答案，那麼我們可以求總數減去不經過黑色的，這樣就是答案了。

那麼也就是求出每個只有紅色邊的連通塊中有幾個點。

總數就是n的k次方，設連通塊有i個點，那麼每次減去i的k次方就行了。

#include <iostream>
#include <cmath>
#include <cstdio>
#include <cstring>
#include <vector>
#include <algorithm>
#include <queue>
#include <set>
#include <map>
#include <stack>
#include <iomanip>
#include <string>

using namespace std;

typedef long long int ll;
ll n,k,sum;
int u,v,x;
const int maxn=100000+5;
vector<int> edge[maxn];
const int mod=1e9+7;
int vis[maxn];

ll quick_m(ll a,ll b){
    ll res=1ll;
    while(b){
        if(b%2) res=res*a%mod;
        a=a*a%mod;
        b/=2;
    }
    return res;
}

void DFS(int x){
    vis[x]=1;
    sum++;
    for(int i=0;i<edge[x].size();i++){
        int nxt=edge[x][i];
        if(!vis[nxt]) DFS(nxt);
    }
}

int main(){
    while(scanf("%I64d%I64d",&n,&k)!=EOF){
        for(int i=0;i<=n;i++) edge[i].clear();
        memset(vis,0,sizeof(vis));
        for(int i=1;i<n;i++){
            scanf("%d%d%d",&u,&v,&x);
            if(x==0){//去掉黑色的邊
                edge[u].push_back(v);
                edge[v].push_back(u);
            }
        }
        ll ans=quick_m(n,k);
        for(int i=1;i<=n;i++){
            if(!vis[i]){
                sum=0;
                DFS(i);
                ans=((ans-quick_m(sum,k))%mod+mod)%mod;
            }
        }
        printf("%I64d\n",ans);
    }
    return 0;
}