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

并查集+快速幂

因为要求的路径至少经过一条黑边，那么反过来思考，我们找出所有只经过白边的路径数量，然后用总共的路径数量减去它即为答案。我们只维护白边，那么一颗树就划分成了若干个连通块，我们用并查集来统计各个连通块的大小然后计数即可。

#include<bits/stdc++.h>
using namespace std;
#define rep(i,a,n) for(int i = a;i<n;i++)
#define per(i,a,n) for(int i = n-1;i>=a;i--)
#define pb push_back
#define mp make_pair
#define eb emplace_back
#define all(x) (x).begin(),(x).end()
#define fi first
#define se second
#define SZ(x) ((int)(x).size())
#define yes cout<<"YES"<<endl;
#define no cout<<"NO"<<endl;
typedef vector<int> VI;
typedef long long ll;
typedef unsigned long long ull;
typedef pair<int,int> PII;
typedef double db;
mt19937 mrand(random_device{}());
const ll mod=1000000007;
int rnd(int x) {return mrand()%x;}
ll gcd(ll a,ll b){return b?gcd(b,a%b):a;}
ll lcm(int a,int b){return a*b/gcd(a,b);}

template<typename T>
struct DSU{
	int n;
	vector<T> p,siz;
	DSU(int n):p(n+1),siz(n+1){
		iota(all(p),0);
		rep(i,0,n+1) siz[i]=1;
	}
	T findd(T x){
		return p[x]==x?x:p[x]=findd(p[x]);
	}
	void unionn(T x,T y){
		x=findd(x),y=findd(y);
		if(x==y) return;
		if(siz[x]>siz[y]) swap(x,y);
		p[x]=y;
		siz[y]+=siz[x];
	}
};

ll qp(ll a,ll b){
	ll res=1;
	while(b){
		if(b&1) res=res*a%mod;
		a=a*a%mod;
		b>>=1;
	}
	return res;
}

void solve(){
	int n,k;
	cin>>n>>k;
	DSU<int> dsu(n);
	rep(i,0,n-1){
		int x,y,c;
		cin>>x>>y>>c;
		if(!c) dsu.unionn(x,y);
	}
	ll ans=qp(n,k);
	rep(i,1,n+1){
		if(dsu.p[i]==i){
			ans-=qp(dsu.siz[i],k);
			if(ans<0) ans+=mod;
		}
	}
	cout<<ans;
}

int main(){
	ios_base::sync_with_stdio(0);cin.tie(0);cout.tie(0);
	//freopen("input.txt","r",stdin);
	//freopen("output.txt","w",stdout);
	int T=1;
	//cin>>T;
	while(T--){
		solve();
	}
	return 0;
}