TopCoder SRM 584 Div1 Medium Excavations

600ptのMedium。本番やった解けた、と思ったけどつめが甘く失敗。
落とした人が多かったところを見ると、pretestが甘めだったかね。
なお、Div2 Hardはこれを簡単にしたもの。
http://community.topcoder.com/stat?c=problem_statement&pm=12641

問題

古代都市にN個の建物が埋まっており、それらの種類kind[i]と埋まっている深さdepth[i]が与えられる。
ここで、N個中K個の建物を選択して、任意の深さDだけ掘ったところ、配列foundに含まれる種類の建物が最低１個以上見つかり、それ以外の種類の建物が見つからなかった。
そのようなN個中K個の建物の選び方を返せ。

解法

まずdepth[i]を座標圧縮しておく。
また、foundに含まれない種類の建物は分類がもはや無意味なので１個にまとめておく。

まず、foundに含まれる種類の建物だけからK個選ぶ場合を列挙する。
この処理では、foundに含まれる種類の建物を最低１個ずつ、合計K個選ぶようにする。
これはDPで処理できる。

次に、foundに含まれない建物を選ぶケースを考える。
foundに含まれない建物のうち、１つを基準として選ぶ。この深さをDとする。
そして、D以上の深さでfoundに含まれない建物をA個選び、D未満の深さでfoundに含まれる建物を最低１個・合計K-(A+1)個選ぶようにする。

自分は本番同じ深さの建物が複数あるケースを見落として失敗した。

int NUM[52][52];
ll comb(int P_,int Q_) {
	static const int N_=102;
	static ll C_[N_][N_];
	if(Q_<0 || P_<Q_) return 0;
	if(C_[0][0]==0) {
		int i,j;
		FOR(i,N_) C_[i][0]=C_[i][i]=1;
		for(i=1;i<N_;i++) for(j=1;j<i;j++) C_[i][j]=(C_[i-1][j-1]+C_[i-1][j]);
	}
	return C_[P_][Q_];
}

class Excavations {
	int N;
	vector<int> D[51];
	int F[51];
	public:
	
	void uniq(vector<int>& DD) {
		vector<int> VV=DD;
		sort(VV.begin(),VV.end());
		VV.erase(unique(VV.begin(),VV.end()),VV.end());
		int i,j,k;
		FOR(i,DD.size()) {
			k=DD[i];
			FOR(j,VV.size()) if(k==VV[j]) DD[i]=j+1;
		}
	}
	
	ll dfs(int depth, int left) {
		vector<pair<int,int> > V;
		ll dp[51][51];
		int i,j,k,l;
		FOR(i,51) if(F[i]) {
			if(NUM[i][depth]==0) return 0;
			V.push_back(make_pair(NUM[i][depth],NUM[i][51]));
		}
		ZERO(dp);
		dp[0][left]=1;
		FOR(i,V.size()) {
			int j=V[i].second;
			for(k=1;k<=j;k++) {
				for(l=k;l<=left-i;l++) {
					// j個中から、k個選ぶが、先頭V[i].firstから１個以上選びたい
					dp[i+1][l-k] += (comb(j,k)-comb(j-V[i].first,k))*dp[i][l];
				}
			}
		}
		return dp[V.size()][0];
	}
	
	long long count(vector <int> kind, vector <int> depth, vector <int> found, int K) {
		int i,j,k;
		map<int,int> cand;
		
		N=kind.size();
		ZERO(F);
		uniq(depth);
		FOR(i,51) D[i].clear();
		ZERO(NUM);
		
		FOR(i,found.size()) F[found[i]]=1;
		FOR(i,N) {
			if(F[kind[i]]) {
				D[kind[i]].push_back(depth[i]);
				NUM[kind[i]][depth[i]]++;
			}
			else {
				cand[depth[i]]++;
				D[0].push_back(depth[i]);
			}
		}
		
		FOR(i,51) FOR(j,51) NUM[i][j+1]+=NUM[i][j];
		FOR(i,51) sort(D[i].begin(),D[i].end());
		ll ret=0;
		// no nofound
		if(K<=N-D[0].size()) ret += dfs(51,K);
		int tot=0;
		ITR(mit,cand) {
			for(i=1;i<=mit->second;i++) {
				for(j=0;j<=min(K-i,(int)D[0].size()-(tot+mit->second));j++) {
					int left=K-(i+j);
					if(left>=found.size() && left<=N-D[0].size()) {
						ll dd=dfs(mit->first-1,left);
						ret += comb(mit->second,i)*comb((int)D[0].size()-(tot+mit->second),j) * dd;
					}
				}
			}
			tot+=mit->second;
		}
		return ret;
	}

};

まとめ

本番中に最後まで詰められなかったのは残念だが、最終的に600ptをノーヒントで解けたのは良かった。
後は本番でここまで詰められるようにしよう。