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

Edgy Trees

time limit per test: 2 second memory limit per test: 256 megabytes input: standard input output: standard output

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 $[a_1,a_2,\dots ,a_k]$ good if it satisfies the following criterion:

• We will walk a path (possibly visiting same edge/vertex multiple times) on the tree, starting from $a_1$ and ending at $a_k$.
• Start at $a_1$, then go to $a_2$ using the shortest path between $a_1$ and $a_2$, then go to $a_3$ in a similar way, and so on, until you travel the shortest path between $a_{k-1}$ and $a_k$.
• 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 $n^k$ sequences of vertices, count how many of them are good. Since this number can be quite large, print it modulo $10^9+7$.

Input

The first line contains two integers $n$ and $k$ ($2\le n\le 10^5$, $2\le k\le 100$), the size of the tree and the length of the vertex sequence.

Each of the next $n-1$ lines contains three integers $u_i$, $v_i$ and $x_i$ ($1\le u_i,v_i\le n$, x i ∈ { 0 , 1 } x_i \in \{0, 1\} xi​∈{0,1}), where $u_i$ and $v_i$ denote the endpoints of the corresponding edge and $x_i$ is the color of this edge ($0$ denotes red edge and $1$ denotes black edge).

Output

Print the number of good sequences modulo $10^9+7$.

Example

$\mathtt{input}$
4 4 1 2 1 2 3 1 3 4 1
$\mathtt{output}$
252

$\mathtt{input}$
4 6 1 2 0 1 3 0 1 4 0
$\mathtt{output}$
0

$\mathtt{input}$
3 5 1 2 1 2 3 0
$\mathtt{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.

Tutorial

由于题目要求中有至少走过一条黑边，所以我们可以用正难则反的思想，求出所有坏序列，即一条黑边也没有，最后再用所有序列减去坏序列即为结果

首先我们可以将黑边删除，剩下的都将是红边连接的分块，对于每个分块，它们自身元素的连接均为坏序列，序列个数为 $s{z}^k$，其中 $zs$​ 为当前分块的节点个数，此时想要建成好序列就需要与其他分块相连，即通过一条黑边

由于所有的序列个数为 $n^k$，所以最终答案为 $n^k-\sum z{s}^k$，其中 $sz$ 为当前分块的节点个数

此解法时间复杂度为 $\mathcal{O}(\alpha (n))$，即并查集的时间复杂度

Solution

#include <bits/stdc++.h>
using namespace std;#define endl '\n'
#define int long long
const int mod = 1e9 + 7; // 998244353;struct DSU { // 并查集vector<int> pre, siz;DSU() {}DSU(int n) {pre.resize(n + 1);std::iota(pre.begin(), pre.end(), 0);siz.assign(n + 1, 1);}int find(int x) {if (pre[x] == x) {return x;}return pre[x] = find(pre[x]);}bool same(int x, int y) {return find(x) == find(y);}bool merge(int x, int y) {x = find(x);y = find(y);if (x == y) {return false;}siz[x] += siz[y];pre[y] = x;return true;}int size(int x) {return siz[find(x)];}
};int ksm(int x, int y, int mod) {x %= mod;int ans = 1;while (y) {if (y & 1) {ans = (ans * x) % mod;}x = (x * x) % mod;y >>= 1;}return ans;
}signed main() {int n, k;cin >> n >> k;int ans = ksm(n, k, mod);DSU dsu(n);vector<int> memo(n + 1);for (int i = 1; i < n; ++i) {int u, v, x;cin >> u >> v >> x;if (not x) {dsu.merge(u, v);}}for (int i = 1; i <= n; ++i) {if (not memo[dsu.find(i)]) {ans -= ksm(dsu.size(i), k, mod);ans = (ans + mod) % mod;memo[dsu.find(i)] = 1;}}cout << ans << endl;return 0;
}