TopCoder SRM 669 Div1 Medium: LineMST

解法

sigma425.hatenablog.com
kmjp.hatenablog.jp

求めるべきMSTは一直線になっているので、ある区間がM未満になる場合を考えるとDPになる。

コード

public class LineMST {
	private final long MOD = (long) 1e9 + 7;
	private long[][] dp;
	private int L;

	public int count(int N, int L) {
		this.L = L;
		dp = new long[N + 1][L + 1];

		// MSTが0, 1, ..., (N-1)を順番につないだものであると仮定する

		// 頂点数Nで最小コストLのものを求める
		long ans = dp(N, L);

		// N!倍して/2する
		for (int i = 3; i <= N; i++) {
			ans = (ans * i) % MOD;
		}
		return (int) ans;
	}

	// 頂点数Nで最小コストMのものを求める
	private long dp(int N, int M) {
		if (N == 1) {
			return 1;
		}
		if (M == 0) {
			return 0;
		}
		if (dp[N][M] > 0) {
			return dp[N][M];
		}

		long cnt = 0;
		for (int k = 1; k < N; k++) {
			// 1からkがM未満のコストの辺で結ばれていて、kとk+1がMのコストの辺で結ばれていて、
			// k+1からNがM以下のコストの辺で結ばれているパターンの数を数える

			long upper = L - (M - 1);// M以上 L以下

			// 1からkまでのk個の頂点と、k+1からNまでの(N-k)個の頂点の間を結ぶ、使われない辺の数。ただしkとk+1を結ぶ辺を除く。
			long unusedEdgeNum = k * (N - k) - 1;
			long pow = modPow(upper, unusedEdgeNum) % MOD;
			cnt += dp(k, M - 1) * dp(N - k, M) % MOD * pow;
			cnt %= MOD;
		}

		// 全ての頂点がM未満のコストで結ばれているパターンも足しておく
		dp[N][M] = (dp(N, M - 1) + cnt) % MOD;
		return dp[N][M];
	}

	private long modPow(long x, long e) {
		long ret = 1;
		long cur = x;
		while (e > 0) {
			if (e % 2 == 1) {
				ret = (ret * cur) % MOD;
			}
			cur = (cur * cur) % MOD;
			e >>= 1;
		}
		return ret;
	}
}