kmjp's blog

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

TopCoder SRM 635 Div1 Hard ColourHolic

SRM

解答まであと１歩だったけど、その１歩が遠い。
http://community.topcoder.com/stat?c=problem_statement&pm=13407

問題

0,1,2,3の４色の商品がそれぞれN[i]個ずつある。
これらをすべて１列に並べるとき、隣同士が同じ色にならない組み合わせは何通りあるか。
なお、N[i]のうち２つは200以下である。

解法

まずNを昇順ソートする。
N[0]もN[1]も0の時は、コーナーケースとして先に処理してしまう。
この時は色順で2,3,2,3,2...か3,2,3,2,3...と交互に並べるしかないので、N[2]とN[3]の差が１以下なら並べられる。

計算量の記載のため便宜上、Nの最大値をP(≦5000)、N[1]の最大値をQ(≦200)と置く。

以降はN[0]かN[1]が1以上であることを前提とする。
まず、色0,1の商品を先に並べ方をDPで求めていく。
その際、同じ色が連続する箇所が何か所あるかも求めておく。
状態数としては[色0の残り商品数][色1の残り商品数][同じ色が連続している箇所の数][最後に置いた商品の色]となりO(Q^3)の状態・計算量で処理できる。


次に、「色2,3の商品をp個のグループに分ける分け方」を考える。
各グループ内は色2と3の商品の数の差が１個以下になるように置かなければならない。
また、商品の数が等しいグループは2と3どちらを先に置くかで2通りずつ置き方がある。

そこで、色3の方が１個多いグループがx個、色3と色2の数が等しいグループがy個、色3の方が１個少ないグループがz個あるとする。
x-y=N[3]-N[2]でなければいけないので、x,y,zのうち２値を総当たりすれば結果的にx,y,zを列挙できる。
色2,3の商品を上記(x,y,z)個のグループに分ける分け方は以下のように考えられる。

  • x+y+z個のグループを上記３通りのグループに分ける方法が {}_{x+y+z} C_{y+z}* {}_{y+z} C_{z}通り
  • yのグループはそれぞれ2通りずつ置き方があるので、 2^y通り
  • 色3の商品はN[3]個あるが、このうちyのグループとzのグループは最低１個は色3の商品を置かなければいけないため、まず１個ずつ割り振りを確定させる。残りleft=N[3]-y-z個はどのグループに割り振るか考えると {}_{left+x+y+z-1} C_{x+y+z-1}

上記３式の積を、「色2,3の商品をp=x+y+z個のグループに分ける分け方DP2[p]」として加算していく。


最後に前半と後半の結果を合わせていく。
色0と1の商品を並べると、商品間と両端で計(N[0]+N[1]+1)個の領域ができるので、そこに色2,3の商品を置いていくことを考える。
各領域は色2,3の商品を置いても置かなくてもいいが、色0,1の商品が連続している箇所は何かしら置かなければならない。

前半のDPにより求めらテル、色0,1を並べた際、同じ色の連続がx個ある並べ方をDP[x]とする。
この場合、色2,3はx以上(N[0]+N[1]+1)グループ以下である分け方をする必要がある。
色2,3の並べ方がyグループ(x≦y≦N[0]+N[1]+1)の場合、yグループ中xグループは色0,1を並べた際、同じ色の連続している箇所に並べなければならず、残り(y-x)グループは空いた(N[0]+N[1]-1-x)グループのどこかに入れればよいので、 DP[x] * \sum_y DP2[y]*{}_{N[0]+N[1]-1-x} C_{y-x}を足していけばよい。

計算量は、O(P+Q^3)かな。PはCombination計算の前処理分。

ll mo=1000000009;

ll combi(ll N_, ll C_) {
	const int NUM_=400001;
	static ll fact[NUM_+1],factr[NUM_+1],inv[NUM_+1];
	if (fact[0]==0) {
		inv[1]=fact[0]=factr[0]=1;
		for (int i=2;i<=NUM_;++i) inv[i] = inv[mo % i] * (mo - mo / i) % mo;
		for (int i=1;i<=NUM_;++i) fact[i]=fact[i-1]*i%mo, factr[i]=factr[i-1]*inv[i]%mo;
	}
	return factr[C_]*fact[N_]%mo*factr[N_-C_]%mo;
}
ll modpow(ll a, ll n) {
	ll r=1;
	while(n) r=r*((n%2)?a:1)%mo,a=a*a%mo,n>>=1;
	return r;
}

int dp[201][201][405][2];
ll dp2[405];

class ColourHolic {
	public:
	
	
	int countSequences(vector <int> n) {
		int i,j,x,y,z;
		sort(n.begin(),n.end());
		if(n[0]==0 && n[1]==0) {
			if(n[2]==n[3]) return 2;
			return (n[2]==n[3]-1);
		}
		
		ZERO(dp);
		ZERO(dp2);
		FOR(x,402) { // plus
			y = x+n[3]-n[2]; // minus
			FOR(z,402) if(x+y+z<=n[0]+n[1]+1 && x+y+z>0) { //same
				int left=n[3]-y-z;
				if(left<0) continue;
				ll tot=combi(x+y+z,x)*combi(y+z,y)%mo*modpow(2,z)%mo*combi(left+(x+y+z-1),(x+y+z-1));
				dp2[x+y+z]=(dp2[x+y+z]+tot)%mo;
			}
		}
		
		if(n[0]>0) dp[n[0]-1][n[1]][0][0]=1;
		if(n[1]>0) dp[n[0]][n[1]-1][0][1]=1;
		for(x=n[0];x>=0;x--) for(y=n[1];y>=0;y--) FOR(i,401) {
			if(dp[x][y][i][0] && x>0) (dp[x-1][y][i+1][0]+=dp[x][y][i][0])%=mo;
			if(dp[x][y][i][0] && y>0) (dp[x][y-1][i][1]+=dp[x][y][i][0])%=mo;
			if(dp[x][y][i][1] && x>0) (dp[x-1][y][i][0]+=dp[x][y][i][1])%=mo;
			if(dp[x][y][i][1] && y>0) (dp[x][y-1][i+1][1]+=dp[x][y][i][1])%=mo;
		}
		ll ret=0;
		FOR(i,402) {
			ll pat=(dp[0][0][i][0]+dp[0][0][i][1])%mo;
			ll t=0;
			for(x=i;x<=401;x++) (t+=dp2[x]*combi(n[0]+n[1]+1-i , x-i))%=mo;
			ret=(ret+pat*t)%mo;
		}
		return (int)ret;
	}
}

まとめ

挿入DPとか、基本的なアイデアは本番にも思いつけたものの本番中に最後まで到達できなかった。