Codeforces Round #309 Div2 C: Kyoya and Colored Balls

問題

codeforces.com

解法

最大の数字kについて考える。kがn個あり、それ以外の数字がm個あるとすると、kのうち1つは一番後ろでなければならない。残りのn-1個の取りうる場所の組み合わせは、k以外の数字の区別をつけなければ _{(m+n-1)} C_{n-1}通りになる。これをkの小さい方から順番に動的計画法で求めていけば良い。

コード

import java.util.Scanner;

public class Main {
	private final int MOD = 1000000007;
	private final int MAX_N = 1002;
	private int K;

	public void solve() {
		int[][] tpas = new int[MAX_N][MAX_N];
		for (int i = 1; i < MAX_N; i++) {
			for (int j = 1; j < MAX_N; j++) {
				if (j == 0) { // 行の左端は 1
					tpas[i][j] = 1;
				} else if (j == i) { // 行の右端も 1
					tpas[i][j] = 1;
					break; // 右端まで計算したので内側のループから抜ける
				} else {
					// 行の途中は一つ上の二つの要素の合計
					tpas[i][j] = tpas[i - 1][j - 1] + tpas[i - 1][j];
					tpas[i][j] %= MOD;
				}
			}
		}

		Scanner sc = new Scanner(System.in);
		K = sc.nextInt();
		if (K == 1) {
			System.out.println(1);
			sc.close();
			return;
		}

		int[] balls = new int[K];
		for (int i = 0; i < K; i++) {
			balls[i] = sc.nextInt();
		}
		sc.close();

		long[] dp = new long[K];
		dp[0] = 1;
		int prevs = balls[0];
		for (int k = 1; k < K; k++) {
			prevs += balls[k] - 1;
			int usable = balls[k] - 1;
			dp[k] = (long) dp[k - 1] * tpas[prevs + 1][usable + 1];
			dp[k] %= MOD;
			prevs++;
		}
		System.out.println(dp[K - 1]);

	}

	public static void main(String[] args) {
		new Main().solve();
	}

}