题解【Codeforces1139C】Edgy Trees

题面

找到所有的只包含红边的连通块，设有 \(n\) 个连通块，每一个连通块中有 \(sz_i\) 个点。

不难发现答案就是 \(\sum\limits_{i=1}^n {sz_i}^k\)。

用并查集维护 \(sz_i\) 即可。

#include <bits/stdc++.h>
#define DEBUG fprintf(stderr, "Passing [%s] line %d\n", __FUNCTION__, __LINE__)
#define File(x) freopen(x".in","r",stdin); freopen(x".out","w",stdout)

using namespace std;

typedef long long LL;
typedef pair <int, int> PII;
typedef pair <int, PII> PIII;

inline int gi()
{
	int f = 1, x = 0; char c = getchar();
	while (c < '0' || c > '9') {if (c == '-') f = -1; c = getchar();}
	while (c >= '0' && c <= '9') x = x * 10 + c - '0', c = getchar();
	return f * x;
}

inline LL gl()
{
	LL f = 1, x = 0; char c = getchar();
	while (c < '0' || c > '9') {if (c == '-') f = -1; c = getchar();}
	while (c >= '0' && c <= '9') x = x * 10 + c - '0', c = getchar();
	return f * x;
}

const int INF = 0x3f3f3f3f, N = 100003, M = N << 1, mod = 1000000007;

int n, m, k;
int tot, head[N], ver[M], nxt[M], edge[N];
int fa[N], sz[N];

inline void add(int u, int v, int w)
{
	ver[++tot] = v, edge[tot] = w, nxt[tot] = head[u], head[u] = tot;
}

int getf(int u)
{
	if (fa[u] == u) return u;
	return fa[u] = getf(fa[u]);
}

map <int, int> p;

inline LL qpow(int x, int y)
{
	LL res = 1;
	while (y)
	{
		if (y & 1) res = (LL)res * x % mod;
		x = (LL)x * x % mod, y >>= 1;
	}
	return res;
}

int main()
{
	//File("");
	n = gi(), k = gi();
	LL ans = qpow(n, k);
	for (int i = 1; i <= n; i+=1) fa[i] = i, sz[i] = 1;
	for (int i = 1; i < n; i+=1)
	{
		int u = gi(), v = gi(), w = gi();
		add(u, v, w), add(v, u, w);
		if (w == 0)
		{
			int x = getf(u), y = getf(v);
			fa[x] = y;
			sz[y] += sz[x];
		}
	}
	for (int i = 1; i <= n; i+=1)
	{
		int u = getf(i);
		if (!p[u])
		{
			p[u] = 1;
			ans = ((ans - qpow(sz[u], k)) % mod + mod) % mod;
		}
	}
	printf("%lld\n", ans % mod);
	return 0;
}