TopCoder SRM 584 Div1 Medium Excavations

問題

n個の遺跡があって、i番目の遺跡は種類kind[i]のものであり深さdepth[i]に埋まっている。


掘削機で遺跡を発掘することができる。
掘削機は掘る深さが一番最初にだけ自由に設定できて(途中で変更はできない)、
Dの深さでi掘ったとき、D≧depth[i]ならその遺跡を発掘することができる。


遺跡をちょうどK個掘った結果、発掘できた遺跡のkindを、
重複しないsetとしてみたときfound[i]であった。
(すなわち、{1, 1, 2, 4, 4}が発掘されたら、foundは{1, 2, 4})


発掘された遺跡の種類の集合がちょうどfoundであるように、掘る遺跡をK個決める、決め方を求めよ。
発掘された遺跡がfoundに一致するならば、Dが異なっても同じものとする。
また、掘る遺跡の集合が同じならば、発掘された遺跡が(重複集合として)異なっていても、同じものとする。

制約条件

n≦50
kind[i]≦50
depth[i]≦10000
found[i]≦50

方針

depthは座標圧縮してよい。


掘る遺跡の集合を決める。
この集合を掘ったとき、発掘されるものがfoundであるようなDで、
適切なものとしてa, bがあったとする。
すると、a≦c≦bを満たすいかなるcであっても、foundは同じになる。


何故かというと、cで発掘されるものはaでも発掘されて、
bで発掘されないものはaでも発掘されないから。


なので、発掘されるものがfoundになるような最小のDがlo, 最大のDがhiであるような
掘る集合の数を数えられれば、掘る集合を漏れなく、重複なく数えられることがわかる。


これどうやって求めたらいいかというと。
これは、次のようなdpをすればよい。


dp[kindをいくつまで見たか][何個掘ったか][最小がちょうどloになるか][最大がちょうどhiになるか] := 場合の数


で、kindごとにそのkindから何個掘るかで遷移する。


最小がちょうどloになるとは、
発掘するkindのうち少なくとも一種類で「ちょうどloの深さの遺跡を発見し、
その種類で掘る残りの遺跡は深さlo以上である」を満たし、
残りの種類で、lo以下の遺跡を一つ以上発掘することが必要十分。


最大がちょうどhiになるとは、
発掘しないkindのうち、少なくとも一種類で「ちょうどhiの深さの遺跡を掘って失敗し、
その種類で掘る残りの遺跡の深さはhi以上である」を満たし、
残りの発掘しない種類で、hiより大の遺跡しか掘らないことが必要十分。


(これはhiがinfの場合だけ例外がある)


で、計算量が一見O(n^5)に見えるんだけど、
制約から、kindが多いと、そのkindにある遺跡は少なくなってしまうので、本当の計算量はたぶんO(n^4.5)とかで、余裕で間に合う。


模範解答はちらっと見たけど何やってるかよくわからなかった。
メモ化再帰らしい。

ソースコード

int n;
int cnt[51][60];

ll dp[2][51][2][2], C[51][51];
bool f[51];

class Excavations {
	public:
	long long count(vector <int> kind, vector <int> depth, vector <int> found, int K) {
		
		rep(i, 51) rep(j, i+1) C[i][j] = j==0||j==i ? 1 : C[i-1][j] + C[i-1][j-1];
		
		memset(cnt, 0, sizeof(cnt));
		memset(f, 0, sizeof(f));
		compress(depth);
		
		n = kind.size();
		rep(i, found.size()) f[found[i]] = 1;
		
		ll ans = 0;
		int D = 0;
		
		rep(i, n){
			cnt[kind[i]][depth[i] + 1]++;
			D = max(D, depth[i] + 1);
		}
		rep(i, 51) rep(j, 55) cnt[i][j + 1] += cnt[i][j];
		
		for(int lo = 1; lo <= D; lo++) for(int hi = lo + 1; hi <= D + 1; hi++){
			
			memset(dp, 0, sizeof(dp));
			dp[0][0][0][0] = 1;
			
			int cur = 0, next = 1;
			
			rep(pos, 51){
				memset(dp[next], 0, sizeof(dp[next]));
				
				rep(i, K + 1) rep(j, 2) rep(k, 2) if(dp[cur][i][j][k]){
				
					ll now = dp[cur][i][j][k];
					
					if(f[pos]) for(int l = 1; l <= cnt[pos][55] && l + i <= K; l++){
						
						//lo未満を1個以上含む
						ll way = C[cnt[pos][55]][l] - C[cnt[pos][55] - cnt[pos][lo - 1]][l];
						dp[next][i + l][j][k] += now * way;
						
						//lo以上でloを含む
						way = C[cnt[pos][55] - cnt[pos][lo - 1]][l] - C[cnt[pos][55] - cnt[pos][lo]][l];
						dp[next][i + l][1][k] += now * way;
					}
					else for(int l = 0; l <= cnt[pos][55] && l + i <= K; l++){
						
						//Rより大で選ぶ
						ll way = C[cnt[pos][55] - cnt[pos][hi]][l];
						dp[next][i + l][j][k] += now * way;
						
						//R以上でRを含む
						way = C[cnt[pos][55] - cnt[pos][hi - 1]][l] - way;
						dp[next][i + l][j][1] += now * way;
					}
				}
				swap(cur, next);
			}
			ans += dp[cur][K][1][1];
			if(hi == D + 1) ans += dp[cur][K][1][0];
		}
		return ans;
	}
	void compress(vi &x){
		vi v;
		rep(i, x.size()) v.pb(x[i]);
		sort(all(v));
		v.erase(unique(all(v)), v.end());
		rep(i, x.size()) x[i] = lower_bound(all(v), x[i]) - v.begin();
	}
};