Codeforces Round 548 (Div. 2) - C Edgy Trees

Link: Edgy Trees

Description:
You are given a tree (a connected undirected graph without cycles) of n vertices. Each of the n−1 edges of the tree is colored in either black or red.

You are also given an integer k. Consider sequences of k vertices. Let's call a sequence [a1,a2,…,ak] good if it satisfies the following criterion:

1. We will walk a path (possibly visiting same edge/vertex multiple times) on the tree, starting from a1 and ending at ak.
2. Start at a1, then go to a2 using the shortest path between a1 and a2, then go to a3 in a similar way, and so on, until you travel the shortest path between ak−1 and ak.
3. If you walked over at least one black edge during this process, then the sequence is good.

Consider the tree on the picture. If k=3 then the following sequences are good: [1,4,7], [5,5,3] and [2,3,7]. The following sequences are not good: [1,4,6], [5,5,5], [3,7,3].

There are nk sequences of vertices, count how many of them are good. Since this number can be quite large, print it modulo 1e9+7.

Input
The first line contains two integers n and k (2≤n≤1e5, 2≤k≤100), the size of the tree and the length of the vertex sequence.

Each of the next n−1 lines contains three integers ui, vi and xi (1≤ui,vi≤n, xi∈{0,1}), where ui and vi denote the endpoints of the corresponding edge and xi is the color of this edge (0 denotes red edge and 1 denotes black edge).

Output
Print the number of good sequences modulo 109+7.

Examples
input
4 4
1 2 1
2 3 1
3 4 1
output
252
input
4 6
1 2 0
1 3 0
1 4 0
output
0
input
3 5
1 2 1
2 3 0
output
210
Note
In the first example, all sequences (4^4) of length 4 except the following are good:

• [1,1,1,1]
• [2,2,2,2]
• [3,3,3,3]
• [4,4,4,4]

In the second example, all edges are red, hence there aren't any good sequences.

Problem solving:
这道题的意思就是给你一个图，图中的边有黑色的有红色的，让你选k个顶点，保证从这k个顶点中的第一个顶点走到最后一个顶点的路径上一定会经过一条黑色的边。并且可以一个点走很多次，问你可以有多少种不同的k个点的组成方式。

显然如果两个点之间所有的边都是红色的，不会满足题意，除了这种情况，都满足。所以我们可以用并查集把每个由红色边组成的连通块中包含的点的个数算出来，总共n个点选k个，总选择方案个数是n^k，减去这些连通块中的满足条件的方案个数就是答案。

Code:

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
const ll    maxn = 1e5 + 10;
const ll    mod = 1e9 + 7;
ll          k, n, f[maxn];
map<ll, ll> ma;
ll find(ll x)
{
    return f[x] == x ? x : f[x] = find(f[x]);
}
void join(ll x, ll y)
{
    x    = find(x), y = find(y);
    f[x] = y; ma[y] += ma[x];   // ma用来存储以y为根节点的连通块中节点的个数
}
ll poww(ll x, ll y)
{
    ll ans = 1;
    while (y)
    {
        if (y & 1)
            ans = (ans * x) % mod;
        x   = (x * x) % mod;
        y >>= 1;
    }
    return ans;
}
int main()
{
    ios::sync_with_stdio(0);
    cin >> n >> k;
    for (ll i = 1; i <= n; i++)
        f[i] = i, ma[i] = 1;
    ll ans = poww(n, k);
    for (ll i = 1, u, v, w; i < n; i++)
    {
        cin >> u >> v >> w;
        if (!w)
            join(u, v);
    }
    for (ll i = 1; i <= n; i++)
        if (f[i] == i)
            ans = (ans - poww(ma[i], k) + mod) % mod;
 
    cout << ans << endl;
    return 0;
}