ABC133 E - Virus Tree 2 (500)

E - Virus Tree 2

概要

N頂点の木を下記の要領で塗り分ける。塗り分けにK色まで使えるとき、塗り分けの通り数を求めよ。(mod 10^9 + 7 して出力)

• 2つの異なる頂点x, yについて、xとyの間の距離が2以下ならxとyは異なる色で塗られる。

制約

1 ≦ N, K ≦ 10^5

方針

0を頂点とする根付き木で考える。Perm(n, k) = n! / (n - k)! とする。
下記のように各頂点の塗り分けの通り数が独立に決まってゆく。

• 根の塗り分け方はK通り
• 根の子がc頂点あるとき、根の子の塗り分け方はPerm(K-1, c)通り
• 根以外の頂点xについてxの子がc頂点あるとき、xの子の塗り分け方はPerm(K-2, c)通り

あとはこれを再帰等で求めればよい。

解答

Submission #6326974 - AtCoder Beginner Contest 133

#include <bits/stdc++.h>
#define rep(i,n) for(int i=(0);i<(n);i++)
 
using namespace std;
 
typedef long long ll;
 
int max_n = 101010;
 
ll MOD = 1e9 + 7;
 
int N, K;
vector<vector<int>> g(max_n, vector<int>());
 
ll dfs(int x, int p){
 
	int num;
	if(p == -1){
		num = K - 1;
	}else{
		num = K - 2;
	}
 
	if(K < g[x].size()){
		return 0;
	}
 
	ll ret = 1LL;
	for(int y : g[x]){
		if(y == p) continue;
 
		ret *= num;
		ret %= MOD;
		num--;
	}
 
	for(int y : g[x]){
		if(y == p) continue;
		ret *= dfs(y, x);
		ret %= MOD;
	}
 
	return ret;
}
 
int main()
{
	cin.tie(0);
	ios::sync_with_stdio(false);
 
	cin >> N >> K;
 
	rep(i, N-1){
		int a, b;
		cin >> a >> b;
		a--;
		b--;
		g[a].push_back(b);
		g[b].push_back(a);
	}
 
	ll ans = K * dfs(0, -1);
	ans %= MOD;
	cout << ans << endl;
}