CF1139C. Edgy Trees 题解 并查集

题目链接：https://codeforces.com/problemset/problem/1139/C

视频讲解：https://www.bilibili.com/video/BV1tZ1FYPELp?p=3

我们可以求总方案数 - 不满足条件的方案数。

设一个不包含黑色边的极大连通块的大小为 \(sz_i\)。则答案为

示例程序：

#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;
}