Codeforces Round #548 (Div. 2) C. Edgy Trees（并查集）

C. Edgy Trees time limit per test2 seconds memory limit per test256 megabytes inputstandard input outputstandard 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 [a1,a2,…,ak] good if it satisfies the following criterion:

We will walk a path (possibly visiting same edge/vertex multiple times) on the tree, starting from a1 and ending at ak.
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.
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 109+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 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 $10^9+7$.

Examples

4 4
1 2 1
2 3 1
3 4 1
252

4 6
1 2 0
1 3 0
1 4 0
0

3 5
1 2 1
2 3 0
210

Note
In the first example, all sequences (44) 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.

题意：

给定一棵树，树中每个节点之间可能为黑色的边或者红色的边，从树中任一节点出发，遍历$k$个节点，至少经过一次黑边的序列为合格序列，求出这样的序列的总个数。

思路：

在长度为$k$的序列里，由于每个位置有$n$种可能，所以总的可能数为$n^k$，用总的可能数减去仅经过红边的序列数即可。由于仅经过红边说明在这些节点在同一个连通区域，并且相互之间只有红边相连。首先统计总的连通区域数，然后利用$n^k-{\sum }_{i=1}^{num}sum[i{]}^k$即为符合条件的序列的个数，其中num是总的连通分量数，sum[i]是每个连通区域的节点数。利用快速幂求得$n^k$,在上面式子求模时，应转化为：$(n^k-{\sum }_{i=1}^{num}sum[i{]}^k+mod)\%mod$。

#include<bits/stdc++.h>
#define int long long
using namespace std;
int f[100005],sum[100005];
vector<int> v;
const int mod=1e9+7;
int findroot(int x){
	if(f[x]==x) return x;
	return f[x]=findroot(f[x]);
}
void merge(int x,int y){
	int x1=findroot(x),y1=findroot(y);
	if(x1!=y1) f[x1]=y1,sum[y1]+=sum[x1];
} 

int cal(int a,int k){
	int ans=1;
	while(k){
		if(k&1) ans=(a%mod*(ans%mod))%mod;
		k=k>>1;
		a=(a%mod*(a%mod))%mod;
	}
	return ans;
}

int n,k;
signed main(){
    ios::sync_with_stdio(false);
    cin.tie(0);
    cin>>n>>k;
    
    for(int i=1;i<=n;i++) f[i]=i,sum[i]=1;
    
    for(int i=0;i<n-1;i++){
    	int u,v,x;
    	cin>>u>>v>>x;
    	if(x==0){
    		merge(u,v);
		}
	}
	for(int i=1;i<=n;i++){
		if(f[i]==i){
			v.push_back(sum[i]);
		}
	}
	
	int ans=cal(n,k);
	
	for(auto num:v){
		ans=(ans-cal(num,k)+mod)%mod;
	}
	cout<<ans<<"\n";
	return 0;
}