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

まさかの１発正答でびっくり。
http://arc029.contest.atcoder.jp/tasks/arc029_4

問題

N頂点からなる根付き木が与えられる。
また、初期状態として各頂点は整数S[i]が格納されている。

ここに、M個の数字T[i]を順次木に挿入できる(してもしなくても良い)。
木への数値の挿入は、任意の頂点に出来る。
ある頂点に数値T[i]を挿入すると、その頂点の数値はT[i]に置き換わり、もともと頂点に入っていた数値は親頂点に挿入される。
同様に、頂点の数値が親方向に１個ずつずれることになる。最終的に根が持っていた数値は捨てられる。

M個の数値を最適に挿入した(または挿入しないことを選択)した場合、木に残った数値の総和の最大値を求めよ。

解法

M個の数値は挿入してもしなくても良いので、小さな数値は挿入する必要がない。
また挿入の位置は選べるので、挿入した数値が以後追い出されないように挿入していくことができる。

よって、M個のうちi個だけ木に挿入する場合、挿入する数値はS[i]の上位i個である。
あとは、「木から数値をi個追い出す場合の追い出された数値の総和の最小値」が求めてあれば、iを総当たりして最大値を求めることができる。

「木から数値をi個追い出す場合の追い出された数値の総和の最小値」は単純な木DPで「各頂点のsubtreeから何個数値を追い出すかの最小値」を求めればよい。
ただし、数値を追い出す場合、１個目はその頂点の数値でなければならない。２個目以降はsubtreeのどれかから選択できる。
この木DPは一見O(N^3)かかるが、subtreeの処理ではループ回数をNでなくsubtreeの頂点数にとどめることで0.5s位で完了することができる。

int N,M;
int S[5001];
int P[5001],num[5001];
vector<int> E[5001];
int T[5001];
ll sum;
ll dpdp[5001][5001];

void solve() {
	int i,j,k,l,r,x,y; string s;
	cin>>N;
	FOR(i,N) cin>>S[i], sum+=S[i];
	FOR(i,N-1) {
		cin>>x>>y;
		x--;y--;
		if(x>y) swap(x,y);
		P[y]=x;
		E[x].push_back(y);
	}
	FOR(x,5001) FOR(y,5001) dpdp[x][y]=1LL<<50;
	for(i=N-1;i>=0;i--) {
		num[i]=1;
		dpdp[i][0]=0;
		dpdp[i][1]=S[i];
		ITR(it,E[i]) {
			for(j=num[i]+num[*it];j>=2;j--) {
				for(k=1;k<=num[*it];k++) {
					if(j-k<1) break;
					dpdp[i][j]=min(dpdp[i][j],dpdp[i][j-k]+dpdp[*it][k]);
				}
			}
			num[i]+=num[*it];
		}
	}
	
	
	cin>>M;
	FOR(i,M) cin>>T[i];
	sort(T,T+M);
	reverse(T,T+M);
	ll ma=sum,ts=0;
	for(i=1;i<=M;i++) {
		ts+=T[i-1];
		ma=max(ma,sum-dpdp[0][i]+ts);
	}
	cout << ma << endl;
}

まとめ

O(N^3)で通るとは思わなかった。