快手2020校园招聘秋招笔试--工程A试卷 题4

题意：中文题面自行理解

思路：处理出每块由红色边连成的连通块的点个数，ans=总方案数-所有只有红边相连的块的总方案数+只有黑边与之相连的点个数

代码：

#include<bits/stdc++.h>
using namespace std;
#define LL long long
const int maxn=1e5+10;
const LL mod=1e9+7;
int n,k;
map<int,int> mp;
vector<int> vec[maxn],tmp;///tmp:所有只有红色边连成的连通块的点个数
int book[maxn];///标记点是否只有黑边与之相连
int fun(int x)
{
    mp[x]=1;
    queue<int> que;
    int ans=1;
    que.push(x);
    while(que.size()){
        int p=que.front();
        que.pop();
        for(int i=0;i<vec[p].size();i++){
            if(mp[vec[p][i]]==0){
                que.push(vec[p][i]);
                mp[vec[p][i]]=1;
                ans++;
            }
        }
    }
    return ans;
}
LL cal(int x)
{
    x=(LL)x;
    LL ans=1ll;
    for(int i=1;i<=k;i++){
        ans=(ans*x)%mod;
    }
    return ans;
}
int main()
{
    scanf("%d%d",&n,&k);
    int u,v,w;
    for(int i=1;i<n;i++){
        scanf("%d%d%d",&u,&v,&w);
        if(w==0){
            vec[u].push_back(v);
            vec[v].push_back(u);
            book[v]=1;
            book[u]=1;
        }
    }
    for(int i=1;i<=n;i++){
        for(int j=0;j<vec[i].size();j++){
            if(mp[vec[i][j]]) continue;
            tmp.push_back(fun(vec[i][j]));
        }
    }
    LL ans=cal(n);
    for(int i=1;i<=n;i++){
        if(book[i]==0){
            ans=(ans-1+mod)%mod;
        }
    }
    for(int i=0;i<tmp.size();i++){
        ans=(ans-cal(tmp[i])+mod)%mod;
    }
    printf("%lld\n",ans);
    return 0;
}