找到所有的只包含红边的连通块,设有 \(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;
}
原文:https://www.cnblogs.com/xsl19/p/12638886.html