TopCoder SRM 564 Div1 Medium AlternateColors2

問題

r個の赤の球、g個の緑の球、b個の青の球がある。
この球を以下の手順を繰り返して並べる。

赤があれば赤の球を並べる。
緑があれば緑の球を並べる。
青があれば青の球を並べる。
（最初に戻る）

r + g + b = nであり、k番目に並べた球は赤い球だったことがわかっている。
n, kが与えられるとき、条件を満たすr, g, bの組の個数を求めよ。

制約条件

n≦100000
k≦n

方針

取り出した球の文字列を考える。
rgbrgbrbrbrrrのようになるが、
この文字列ひとつについて、数の組(r, g, b)が一対一で対応している。

なので、この文字列のうち、条件を満たすものが何通りあるかを数えればよい。

dp[i][j][k]を、今球をi個並べて、直前がjの色で、集合kの球を使い切った状態であるようなときの場合の数とする。
これを更新していくdpをすればいい。

場合わけして数式で直接的に求めることもできる。

ソースコード

ll dp[100010][8][3];

class AlternateColors2 {
	public:
	long long countWays(int n, int k) {
		
		memset(dp, 0, sizeof(dp));
		for(int i = 1; i < 8; i++) dp[0][i][2] = 1;
		
		rep(i, n) rep(j, 8) rep(p, 3){
			if(dp[i][j][p] == 0) continue;
			if(j == 0) continue;
			
			int nxt = p;
			rep(l, 3){
				nxt = (nxt + 1) % 3;
				if(j & 1 << nxt) break;
			}
			if(i == k - 1 && nxt != 0) continue;
			
			dp[i + 1][j][nxt] += dp[i][j][p];
			dp[i + 1][j ^ 1 << nxt][nxt] += dp[i][j][p];
		}
		ll ans = 0;
		rep(i, 3) ans += dp[n][0][i];
		return ans;
	}
};