并查集加数学

题目:https://codeforces.com/contest/1139/problem/C

#include<iostream>
#include<map>
using namespace std;
typedef long long ll;
const ll maxn = 1e5 + 10;
const ll mod = 1e9 + 7;
ll f[maxn];
ll n, k;
//n从2开始
void init() { for (int i = 1; i <= n; ++i) { f[i] = i; } }//自己的祖先是自己
ll find(ll x)
{
    if (x == f[x]) { return f[x]; }
    else { return f[x]= find(f[x]); }
}
void Merge(ll x, ll y)
{
    x = find(x); y = find(y);
    if (x == y) { return; }
    else { f[x] = y; }
}
ll quickpow(ll x, ll n)
{
    ll ans = 1;
    while (n)
    {
        if (n & 1) { ans = (ans * x) % mod; }
        x = (x * x) % mod;
        n >>= 1;
    }
    return ans % mod;
}
int main()
{
    map<ll, ll>mp;
    int n, k;
    cin >> n >> k;
    int size, len,color;
    init();
    for (int i = 1; i <= n - 1; ++i)
    {
        cin >> size >> len >> color;
        if (!color) { Merge(size, len); }//红色的集合进行合并
    }
    for (int i = 1; i <= n; ++i) { mp[find(i)]++; }
    ll ans = quickpow(n, k);//总集合数
    //总集合减去相应集合
    for (auto it=mp.begin();it!=mp.end();++it)
    {
        ans = (ans + mod - quickpow(it->second, k)) % mod;
    }
    cout << ans << endl;
    return 0;

}