TopCoder SRM 564 Div1 Medium AlternateColors2

続いてDiv1。
本番はかなり手こずったけど、最終的に解けて良かった。
解き方はまだしも、もう少しきれいにすんなり書きたかったけどね。
http://community.topcoder.com/stat?c=problem_statement&pm=12345

Div2と同じようにボールを減らすが、今回の場合K個目が赤になるような合計N個の赤緑青の組み合わせの数を答えるもの。
自分は場合分けして組み合わせを数えてみた。以下を組み合わせればよい。

• K%3==1の場合、K個目までは赤緑青がそろっている状態。残りのN-K個を赤緑青で分け合うので、通り
• K個目までに青がなくなり、赤と緑が残っている場合。青の数ごとに、K個目までに減らす赤と緑が決まるので、残りのN-K個を赤緑で分ければよい。通り。
• 青と緑が逆の場合も同じ組み合わせ。
• 青も緑もK個目までになくなった場合。赤の数ごとに、青と緑が(赤-2)個より少なくなる組み合わせ。

class AlternateColors2 {
	public:
	long long countWays(int n, int k) {
		ll res=0;
		int r,g,b,gb;
		ll l;
		
		// gもhもあまる
		if((k%3)==1) {
			r=k/3+1;
			g=k/3;
			b=k/3;
			l = n-k;
			// l個をrとgとhで割る (l+2)_C_2
			res += (l+2)*(l+1)/2;
		}
		
		// gが充足、bが不足 (およびその逆)
		for(b=0;b<k;b++) {
			if((k-b)%2==0) continue;
			r = ((k-b)+1)/2;
			g = ((k-b)-1)/2;
			if(b>=r-1) continue;
			l = n-k;
			// l個をrとgで割る 2*((l+1)_C_1)
			res += 2*(l+1);
		}
		
		// gもhも不足
		for(r=0;r<=k;r++) {
			if(r+(r-1)+(r-1)<k) continue;
			if(k-r<=r-2) {
				//k-r個を分ける
				res += k-r+1;
			}
			else {
				// k-r個を分けるが上限はr-2
				int mx=min(k-r,r-2);
				int mn=k-r-mx;
				if(mx>=mn) res += (mx-mn)+1;
			}
		}
		
		return res;
	}
};

別解もあるみたいね。
赤緑青の残る・残らないで計８通りに対し、DPで解いてもよいみたい。
終了直後のwriterの反応を見ると、DPは想定してなかったようね。

まとめ

シンプルだけど面白い問題。
数え上げよりもDPの方がスマートに見えるな…。
ともかく解けて良かった。