Codeforces #240 Div1 D. Mashmokh and Water Tanks

問題設定が妙にわかりにくいな…。
http://codeforces.com/contest/414/problem/D

問題

M頂点の根付き木と、数値k,pが与えられ、以下のゲームを行う。

k個以内の異なる頂点を選択し、その頂点にスコア1を与える。他の頂点はスコア0である。
ここから各頂点のスコアが0になるまで以下の処理を行う。

• 毎ターン、幾つかの頂点を選択してclose状態にすることができる。この際、1ターンclose状態にするのにその頂点のスコア分のコストがかかる。
  • ただし、root頂点はcloseに出来ない。
• close状態の選択後、closeでない各頂点のスコアは1つ親の頂点に移動する。
  • 親頂点がcloseでも移動ができる。
  • root頂点分のスコアはグラフ外に出ていく。

各ターンでグラフ外に出たスコアの最大値がゲームのスコアとなる。
頂点をcloseにするコストの上限がpである時、得られる最大スコアを求めよ。

解法

まずDFSで各頂点の深さを求めておく。

スコアを最大化するには、1つのターンにできるだけまとめてスコアを排出すればよい。
頂点のcloseを行わなければ、その頂点の深さ分のターン後にスコアが排出される。

そこで排出予定ターンを総当たりし、得られるスコアを求めればよい。
排出予定ターンtに対し、深さtの頂点がk個以上あれば、それらに初期スコアを与えればtターン後にスコアkが得られる。
深さtの頂点がk個未満なら、深さt-1、深さt-2…の頂点に順に割り当てていけばよい。
その際、深さのtとの差の分だけコストがかかるので、その総和をp以下となる範囲でできるだけおいていく。

事前に部分和を取っていくと、選んだ頂点数に対するコストがO(1)で求められるので二分探索でコストp以下の頂点数を求めればよい。

排出予定ターンがO(M)、二分探索でO(logM)なので全体の計算量はO(M logM)。

int M,K,P;
vector<int> E[100001];
int D[100001],DN[100001],MD;
vector<int> VD;
ll C[100002];


void dfs(int cur,int pre) {
	int i;
	if(pre!=-1) D[cur]=D[pre]+1;
	MD=max(MD,D[cur]+1);
	DN[D[cur]]++;
	FOR(i,E[cur].size()) if(E[cur][i]!=pre) dfs(E[cur][i],cur);
}

void solve() {
	int f,i,j,k,l,x,y;
	
	cin>>M>>K>>P;
	FOR(i,M-1) {
		cin>>x>>y;
		E[x-1].push_back(y-1);
		E[y-1].push_back(x-1);
	}
	dfs(0,-1);
	for(i=1;i<MD;i++) FOR(j,DN[i]) VD.push_back(i);
	for(i=VD.size()-1;i>=0;i--) C[i]=C[i+1]+100000LL-VD[i];
	
	int ma=min(1,K);
	FOR(i,VD.size()) {
		if(i<VD.size()-1 && VD[i]==VD[i+1]) continue;
		int L=max(0,i-(K-1)),R=i;
		while(L<R) {
			ll po=(L+R)/2;
			ll cost=C[po]-C[i+1]-(100000LL-VD[i])*(i+1-po);
			if(cost<=P) R=po;
			else L=po+1;
		}
		ma=max(ma,i-L+1);
	}
	
	cout << ma << endl;
}

まとめ

1500ptとはいえD問題が1発正答できてよかった。