ARC 029 D - 高橋君と木のおもちゃ

問題

n頂点からなる木がある。木には最初s[i]の値がかかれている。
次の操作をm回行う。

数字tiを手に入れる。捨てるか、好きな木の頂点cを選んで数字を置く。
頂点cにもともとあった数字はcの親に行く。cの親vにあった数字はvの親に行く。
……と同様に根までこの操作が続けられ、根にもともとあった数字は消える。

最後に木に残っている数字の和を最大化したい。
和の最大値を求めよ。

制約条件

n, m≦5000

方針

よくある木で配列を持ちまわすDP.

結果的にi個、元の数字を消すとするとき消す頂点の集合というのは根を含み、連結なってなければならない。
したがって根から消す頂点i個を、和が最小になりかつ連結になるよう選べばいいことがわかる。

愚直にdp[v][i]を、vの頂点以下i個を消すときの和の最小値というDPだとO(n^3)でTLE.
ここでdp[i]の配列を持ちまわすようなDPをするとO(n^2)でできる。

rec(頂点v, 親から持ってきたdp配列)という感じで再帰関数を書くことを考える。
まず、子にdp配列を投げるとき。
子に1つでも消す頂点があるときは、vも消さなければならないので、
子に投げる配列では必ずvを使うようにしなければならない。
で、配列を親に返すとき、各iで最小値のほうを採用すればよい。

なんかデータが弱くてO(n^3)でも通るらしい。
[追記]
データが弱いわけではなくて、自明なDPでもO(n^2)になるのだった。参考→http://topcoder.g.hatena.ne.jp/iwiwi/20120428/1335635594

T(l + r) = T(l) + T(r) + l * rの計算量はT(n) = O(n^2)。
T(n)≦cn^2とおいて帰納法すればよくて、T(l + r) = cl^2 + cr^2 + l*r ≦ c(l+r)^2 ≦ cn^2

ソースコード

int n, s[5000], m, t[5000];
vector<vi> e;

void rec(int c, ll *a){
	ll *b = new ll[n + 1];
	b[0] = inf;
	rep(i, n) b[i + 1] = a[i] + s[c];
	
	rep(i, e[c].size()) rec(e[c][i], b);
	rep(i, n + 1) a[i] = min(a[i], b[i]);
	
	delete [] b;
}

int main(){
	cin >> n;
	e.resize(n);
	rep(i, n) cin >> s[i];
	rep(i, n - 1){
		int a, b; cin >> a >> b;
		a--; b--;
		e[a].pb(b);
	}
	cin >> m;
	rep(i, m) cin >> t[i];
	sort(t, t + m, greater<int>());
	
	ll *a = new ll[n + 1];
	rep(i, n + 1) a[i] = inf; a[0] = 0;
	rec(0, a);
	
	ll ans = 0, sum = accumulate(s, s + n, 0ll);
	rep(i, n + 1){
		ans = max(ans, sum - a[i]);
		sum += t[i];
	}
	cout << ans << endl;
	
	return 0;
}