CF1139C. Edgy Trees 题解 并查集
题目链接：https://codeforces.com/problemset/problem/1139/C
视频讲解：https://www.bilibili.com/video/BV1tZ1FYPELp?p=3
我们可以求总方案数 - 不满足条件的方案数。
设一个不包含黑色边的极大连通块的大小为 \(sz_i\)。则答案为
\[n^k - \sum \{ sz_i^k \}
\]
示例程序：
#include <bits/stdc++.h>
using namespace std;
const int maxn = 1e5 + 5;
const long long mod = 1e9 + 7;
int n, k, f[maxn], sz[maxn];
void init() {
for (int i = 1; i <= n; i++)
f[i] = i, sz[i] = 1;
}
int find(int x) {
return x == f[x] ? x : f[x] = find(f[x]);
}
void funion(int x, int y) {
int a = find(x), b = find(y);
if (a != b) {
f[b] = a;
sz[a] += sz[b];
}
}
long long mypow(long long a, long long n) {
long long res = 1;
for (int i = 0; i < n; i++)
res = res * a % mod;
return res;
}
int main() {
scanf("%d%d", &n, &k);
init();
for (int i = 1; i < n; i++) {
int a, b, c;
scanf("%d%d%d", &a, &b, &c);
if (!c)
funion(a, b);
}
long long ans = mypow(n, k);
for (int i = 1; i <= n; i++)
if (i == find(i))
ans = (ans - mypow(sz[i], k) + mod) % mod;
printf("%lld\n", ans);
return 0;
}