CF1139C Edgy Trees

题目地址：CF1139C Edgy Trees

红黑树

\(ans\) 应该等于总数（\(n^k\)）减去不含黑色边的序列数量

不含黑色边就意味着一个序列只能在一个红色联通块中

一个红色联通块中的序列数量应该是点数的 \(k\) 次方

求联通块用dfs用并查集都可以

然后快速幂一下再一减就是 \(ans\)

#include <bits/stdc++.h>
#define ll long long
using namespace std;
const int N = 1e5 + 6, P = 1e9 + 7;
int n, k, v[N], cnt;
vector<int> e[N];

inline int ksm(int a, int b) {
    int c = 1;
    while (b) {
        if (b & 1) c = (ll)c * a % P;
        a = (ll)a * a % P;
        b >>= 1;
    }
    return c;
}

void dfs(int x) {
    v[x] = 1;
    ++cnt;
    for (unsigned int i = 0; i < e[x].size(); i++) {
        int y = e[x][i];
        if (v[y]) continue;
        dfs(y);
    }
}

int main() {
    cin >> n >> k;
    for (int i = 1; i < n; i++) {
        int x, y, z;
        scanf("%d %d %d", &x, &y, &z);
        if (z) continue;
        e[x].push_back(y);
        e[y].push_back(x);
    }
    int ans = ksm(n, k);
    for (int i = 1; i <= n; i++)
        if (!v[i]) {
            cnt = 0;
            dfs(i);
            ans = (ans - ksm(cnt, k) + P) % P;
        }
    cout << ans << endl;
    return 0;
}