Codeforces 414(#240 Div1) D. Mashmokh and Water Tanks

問題

n頂点からなる木が与えられる。根は0番のノード。
各頂点には水槽がついている。お金をp円もっている。

最初にk個以下の好きなノードを選び、そのノードの水槽に1リットルずつ水を入れる。
次から全ての水槽が空になるまで以下の操作を繰り返す。

• 水槽のふたを全部あける
• 好きな水槽を（複数でもよい）選び、ふたを閉める。wリットル水が入った水槽のふたを閉めるのにかかる料金はw円
• 根の水槽を空にする
• 木の深さ1の頂点の水槽の水を、その親の水槽に移し変える。その水槽のふたが閉まっていたら飛ばす。親のふたは閉まっていても関係ない。

一つの水槽に（操作の途中の状態も含む）入っていた水の量の最大値を求めよ。

制約条件

n≦10^5
k, p≦10^9

方針

最大値をWとすると、どこかの頂点にWの水が入るタイミングがあれば、必ず根にWの量の水が入る時間があるので、
根に時間tに全ての水が集まると考えてよい。

時間tには深さtまでの頂点の水が来る。
深さiの頂点の水１リットルを根にとどめるためには、(t-i)円かかる。

なので、tを決めたら、深い順に貪欲に頂点を使っていけばよい。
tを1深くしたときのコストの差分は、ノードの個数だけで簡単に計算できるので、
あらかじめ深さiにある頂点の個数を求めておいて、
しゃくとり法を使えばO(n)で計算できる。

ソースコード

int n, lim_node, lim_cost;
vector<vi> e;

int cnt[100100];

void rec(int c, int p = 0, int d = 0){
	cnt[d]++;
	rep(i, e[c].size()) if(e[c][i] != p) rec(e[c][i], c, d + 1);
}

int main(){
	scanf("%d%d%d", &n, &lim_node, &lim_cost);
	e.resize(n);
	rep(i, n - 1){
		int a, b;
		scanf("%d%d", &a, &b); a--; b--;
		e[a].pb(b);
		e[b].pb(a);
	}
	rec(0);
	cnt[0] = 0;
	
	int ans = 1, nodes = 0;
	ll cost = 0;
	
	for(int i = n, j = n; j > 0; j--){
		
		while(i > 0){
			if(cost + cnt[i] * (ll)(j - i) > lim_cost) break;
			
			cost += cnt[i] * (ll)(j - i);
			nodes += cnt[i--];
		}
		
		int rem = (lim_cost - cost) / (j - i);
		ans = max(ans, nodes + min(rem, cnt[i]));
		
		nodes -= cnt[j];
		cost -= nodes;
	}
	ans = min(ans, lim_node);
	cout << ans << endl;
	
	return 0;
}