TopCoder SRM 564 Div1 Medium: AlternateColors2

解法

ekaing.hatenablog.com

このページの別解がすごすぎた。

コード

public class AlternateColors2 {

	public long countWays(int n, int k) {
		long ans = 0;
		for (int red = 1; red <= k; red++) {
			int notred = k - red;
			int step = red - 1;

			if (step * 2 < notred) {
				continue;
			}

			if (step * 2 == notred) {
				// k番目のボールが破壊されるまで全てのステップが実行された場合
				int rest = n - k;
				// restこのボールはRGBどれでも良い
				ans += (long) (rest + 2) * (rest + 1) / 2;
			} else if (step > notred) {
				// k番目のボールが破壊されるまでに確実に全てのGBボールが破壊される場合
				ans += notred + 1;
			} else {
				// k番目のボールが破壊されるまでに確実に全てのGBボールが破壊される場合
				ans += step - (notred - step) + 1;
				// GBどちらかが残るパターン 各n-kパターン
				ans += (n - k) * 2;
			}
		}

		return ans;
	}
}