TopCoder SRM 659 Div1 Medium CampLunch

豹本に似た問題があったね。
http://community.topcoder.com/stat?c=problem_statement&pm=13713

問題

M人の人がN回食事をとる。
食事は毎回M人が1列に並んで食事をする。

各人に与えられる食事のパターンは以下のいずれかである。

• 1人分の食事を与えられ、1回の食事で食べきる。
• 2人分の食事を与えられ、1人で2回分の食事で食べきる。
• 2人分の食事を与えられ、隣の席の2人で一緒に1回の食事で食べきる。

食事は毎回過不足なく全員に行きわたるとする。
そのような全体の食事のパターン数を求めよ。

解法

豹本p.349、床板の問題のアイデアを借用できる。
すなわち、y回目の食事で2回分の食事を与えられてしまった人をbitmapで管理し、その人は(y+1)日目には食事を与えられない、と考えてDPで解ける。
ただ、この場合O(N*2^2M)ほどかかってしまいTLEする。

ここで高速ゼータ変換の要領で計算量をO(N*M*2^M)に落とすことを考える。
y回目に取った食事が1人前か2人前(2人で2人前食べたか、(y-1)回とy回で2回分1人で食べたか)かどうかをM人分bitmaskで表し、1～y回目までの食事の組み合わせ数をdp[y][mask]とする。

y+1回目の食事の組み合わせを考える際、まずは2人隣同士の人で食事を分けるパターンを無視し、y回目と(y+1)回目で1人が2人分の食事を食べきるケースを考える。
y回目と(y+1)回目で1人が2人分の食事を食べきるためには、y回目の時点だけみれば1人分の食事を取っている場合、(y+1)回目と合わせて2人分の食事をとれる。
よってdp[y+1][mask]は以下のように計算できる。

上記計算を各mask,mについて行うとO(2^2M)かかりTLEする。
しかしここで高速ゼータ変換の要領で計算するとO(M*2^M)に抑えられる。

その後、隣の席の2人で食事をとるケースについて、同様にO(M*2^M)で値を更新していく。
すなわちa番目の人とb番目の人が隣接している場合、(mask & (2^a | 2^b))==0であるmaskに対して

で計算できる。

ll mo=1000000007;
ll dp[1<<16];

class CampLunch {
	public:
	
	int count(int N, int M, vector <string> a) {
		int i,x,y,mask,mask2;
		
		ZERO(dp);
		dp[0]=1;
		
		FOR(y,N) {
			FOR(x,M) {
				FOR(mask,1<<M) if((mask & (1<<x))==0) (dp[mask]+=dp[mask^ (1<<x)])%=mo;
			}
			
			FOR(x,M-1) {
				int fil =(1<<(a[y][x]-'A')) | (1<<(a[y][x+1]-'A'));
				FOR(mask,1<<M) if((mask&fil)==0) (dp[mask | fil] += dp[mask])%=mo;
			}
			
			FOR(mask,1<<M) {
				mask2 = ((1<<M)-1)^mask;
				if(mask<mask2) swap(dp[mask],dp[mask2]);
			}
		}
		ll ret=0;
		FOR(mask,1<<M) ret += dp[mask];
		
		return ret%mo;
	}
}

まとめ

高速ゼータ変換、なかなかさらっと書けないんだよなぁ。