TopCoder SRM 564 Div1 Medium: AlternateColors2

問題

TopCoder Statistics - Problem Statement

コード

public class AlternateColors2 {

    public long countWays(int n, int k) {
	long ans = 0;
	for (int red = 1; red <= k; red++) {
	    int notred = k - red;// k番目の赤ボールが破壊される前に破壊される赤以外のボールの数
	    int step = red - 1;

	    if (step * 2 < notred) {
		// notredがstep*2以内に収まらない場合、
		// k番目に破壊されるはずのボールが破壊されなくなってしまう
		continue;
	    }

	    if (step * 2 == notred) {
		// 全てのステップでRGBが破壊されている場合
		long rest = n - k;
		// k番目に赤が破壊された後のn-k個のボールの色は何でも良い。
		// よってこの時、重複組み合わせで(rest+2)!/(rest!2!)
		ans += (rest + 2) * (rest + 1) / 2;
	    } else if (notred < step) {
		// k番目に赤が破壊される前に全てのBGが破壊されてしまう時
		// 考えうる青の個数は0〜notredの notred+1 通りで、赤の個数、緑の個数はそれによって1通りに確定する
		ans += notred + 1;
	    } else {
		// k番目に赤が破壊される前に全てのBGが破壊される場合、
		// 考えうる青の個数は(notred-step)〜stepの step-(notred-step)+1 通りである
		// 赤の個数、緑の個数はそれによって1通りに確定する
		ans += step - (notred - step) + 1;

		// k番目の赤が破壊された後、青が残っているとすると、青の個数はred個以上である。この個数はn-k通りである。
		// この時、緑は残っていないので、緑の個数はnotred-step個の1通りである。
		ans += n - k;

		// k番目の赤が破壊された後、緑が残っているとした場合も同様
		ans += n - k;
	    }
	}
	return ans;
    }
}