木をK色で塗り分ける場合の数 - Senの競技プログラミング備忘録

つぎのような問題を考える。

サイズNの無向木が与えられる。木の頂点に色を塗っていき、隣あった頂点には同じ色を塗ってはならない。
木をちょうどK色で塗るとき、何通りの塗り分け方があるか？ $10^9+7$ で割ったあまりを答えよ。

制約：　2 <= K <= N <= 1000

例えば次のような問題。
codeforces.com

似た考えでとける問題。ある意味これよりもさらに易しい。
atcoder.jp

考察

木の上をちょうどK色で塗るという条件は、考えてみるとわかるがかなりシビアである。木の上でDPを考えるにしても、今まで何種類の色を使って、次に使う色は新しい種類か、それとも昔使った種類か？隣と同じ色になるか？などと、大変煩雑である。

Point
「Nぴったりの○○」が考えづらいのなら、「N以下のすべての○○」で考えて、余分なのを引く策を考える！

今回の問題はまさにこの考えを使うド典型問題であると思う。
ちょうどK色で塗るというには難しいが、「K色以下で塗る」場合の数ならば、話は別である。

木の1つの頂点を選び、それを根とする。

根の塗り方は、K色どれで塗ってもよいので、K通り
その子供の塗り方は、「親と違う色で塗る」と考えて塗れば、実は木の形と関係なく、親と同じ色以外のK - 1通りで塗れる。

親頂点から見て、子頂点を塗る際に、孫頂点の数などに関しては、親から子へという順序考えれば、孫の頂点の塗り方は、子の頂点の塗り方を決めた時に初めて考えるものなので、子の頂点の塗れる色の数を考える際に考慮しなくていいとわかる。

これらの考察から、場合の数はNとKのみに依存し、 $K(K-1)^{n-1}$ 通りであるとわかる。(先ほど述べた考察の○○通りの部分はすべて独立してるため、累乗で場合の数が書ける。)

さて、ここから不要な部分を引いていこう。
「K色以下で木を塗る」というのは「色1から色Kまで使用可能の時に木を塗る」ということである。
ここで注意すべきなのは、「色1から色Kまで使用可能の時に木を塗る」から「色1から色K - 1まで使用可能の時に木を塗る」を引いた値が答えではないということである。これを引いても、ちょうどK色使って塗る場合の数が得られるわけではない。(例えば、得られるのは色1, 色Kの2色のみで塗られた場合も数えられてしまう。得られるのは、色Kを使った塗り分けの場合の数である)

本当に引くべきなのは、「ちょうどi(2 <= i <= K - 1)色で木を塗り分ける場合の数」× ${}_K \mathrm{C} _i$
である。

式の意味を解説すると、ちょうどi色で木を塗り分ける場合の数がわかっているのなら、その使われてるi色をK色の中から選ぶという意味で ${}_K \mathrm{C} _i$ をかけてる。そして、これを2 <= i <= K - 1まで引けば、残るのはちょうどK色を使った塗り分けの場合の数である。

さて、この式では、「ちょうどK色で木を塗り分ける場合の数」を計算するために、「ちょうどi(2 <= i <= K - 1)色で木を塗り分ける場合の数」が必要としてるため、漸化式を立てて動的計画法で計算することができる。この式で表すメリットは、まさにこれである。

よって、動的計画法で「ちょうどi色で木を塗り分ける場合の数」を計算しながらiを増やしていくことで、 $O(N(K+ \log N))$ で解くことができる。( $\log N$ は繰り返し二乗法)

ACソースコード：
ライブラリのModintは省略した。mintがModintである。

#include<iostream>

using namespace std;

typedef long long ll;

namespace comb {
	ll fact[500001];
	ll fact_inv[500001];
	const ll MOD = 1000000007;

	ll modpow(ll a, ll b) {
		ll base = 1, kakeru = a;
		while (b > 0) {
			if (b & 1)base *= kakeru, base %= MOD;
			kakeru *= kakeru, kakeru %= MOD;
			b >>= 1;
		}
		return base;
	}

	void setup() {
		fact[0] = 1, fact[1] = 1;
		for (int i = 2; i <= 500000; i++) {
			fact[i] = (fact[i - 1] * i) % MOD;
		}

		fact_inv[500000] = modpow(fact[500000], MOD - 2);

		for (int i = 499999; i >= 0; i--) {
			fact_inv[i] = fact_inv[i + 1] * (i + 1);
			fact_inv[i] %= MOD;
		}

	}

	ll combination(int a, int b) {
		if (a - b < 0)return 0;
		ll ret = fact[a];
		ret *= fact_inv[b], ret %= MOD;
		ret *= fact_inv[a - b], ret %= MOD;
		return ret;
	}

	ll permutation(int a, int b) {
		if (a - b < 0)return 0;
		ll ret = fact[a];
		ret *= fact_inv[a - b], ret %= MOD;
		return ret;
	}

	ll homogeneous(int a, int b) {
		return combination(a + b - 1, b);
	}

}

int n, k;
mint dp[1000 + 10];//dp[i] := ちょうどi色で塗り分けた時の場合の数

int main() {
	comb::setup();
	cin >> n >> k;
	for (int i = 0; i < n - 1; i++) {
		int parent;
		cin >> parent;
	}

	mint ans = 0;
	for (int i = 2; i <= k; i++) {
		mint tmp = i * modpow(i - 1, n - 1);
		//色1 ~ iのみを使って、ぬり分けた時の場合の数
		for (int j = 2; j <= i - 1; j++) {
			tmp -= mint(comb::combination(i, j)) * dp[j];
		}
		dp[i] = tmp;
	}
	cout << dp[k] << endl;
	return 0;
}