kmjp's blog

競技プログラミング参加記です

TopCoder SRM 551 Div2 Hard ColorfulCupcakesDivTwo

Div2はHardでも950pt以下の問題はだいぶ易しめ。
http://community.topcoder.com/stat?c=problem_statement&pm=12138

問題

3色(A,B,C)のケーキが計N個与えられる。
円形のテーブルに並んだN人にケーキを配る際、隣同士が同じ色のケーキを持たないように配る組み合わせ数を答えよ。

解法

隣接する人が同じケーキを持ってはいけないので、状態としてdp[Aの残り数][Bの残り数][Cの残り数][直前の色]を持ってDPを行えばよい。
問題は人の並びが円形であることだが、それは1番目の人の色をA,B,C3通り置いて3回ループさせ、N番目が1番目の色と重ならないような組み合わせ数を数えるだけで良い。
(DPの状態として1番目の人の色を持たせても良いが、結局3倍計算することに変わりはない。)

ll dp[2][51][51][51][3];
ll mo=1000000007;

class ColorfulCupcakesDivTwo {
	public:
	int countArrangements(string cupcakes) {
		int c[3],i,L,x,y,z,col;
		ll ret = 0;
		
		L=cupcakes.size();
		ZERO(c);
		FOR(i,L) c[cupcakes[i]-'A']++;
		ZERO(dp);
		
		FOR(col,3) {
			if(c[col]==0) continue;
			dp[1][c[0]-(col==0)][c[1]-(col==1)][c[2]-(col==2)][col]=1;
		
			for(i=1;i<L;i++) {
				int cur=i%2, tar=cur^1;
				ZERO(dp[tar]);
				FOR(x,L) FOR(y,L) {
					z = L-i-x-y;
					if(z<0) break;
					
					if(dp[cur][x][y][z][0]) {
						if(y) dp[tar][x][y-1][z][1] = (dp[tar][x][y-1][z][1] + dp[cur][x][y][z][0]) % mo;
						if(z) dp[tar][x][y][z-1][2] = (dp[tar][x][y][z-1][2] + dp[cur][x][y][z][0]) % mo;
					}
					if(dp[cur][x][y][z][1]) {
						if(x) dp[tar][x-1][y][z][0] = (dp[tar][x-1][y][z][0] + dp[cur][x][y][z][1]) % mo;
						if(z) dp[tar][x][y][z-1][2] = (dp[tar][x][y][z-1][2] + dp[cur][x][y][z][1]) % mo;
					}
					if(dp[cur][x][y][z][2]) {
						if(x) dp[tar][x-1][y][z][0] = (dp[tar][x-1][y][z][0] + dp[cur][x][y][z][2]) % mo;
						if(y) dp[tar][x][y-1][z][1] = (dp[tar][x][y-1][z][1] + dp[cur][x][y][z][2]) % mo;
					}
				}
			}
			
			ret += dp[L%2][0][0][0][0]+dp[L%2][0][0][0][1]+dp[L%2][0][0][0][2];
			ret -= dp[L%2][0][0][0][col];
		}
		
		return ret % mo;
	}
};

まとめ

Div2とはいえずいぶんひねりのない問題。