EPIC Institute of Technology Round Summer 2024 : E. Wonderful Tree!

配点の割に妙に時間がかかってしまった。
https://codeforces.com/contest/1987/problem/E

問題

根付き木が与えられる。
その各頂点vには整数A[v]が設定されている。

各点vのA[v]は、vの子頂点cにおけるA[c]の総和以下とならなければいけない。
各点のAの値を任意に増加できるとき、条件を満たすための総増加量の最小値を求めよ。

解法

A[v]がvの子頂点cにおけるA[c]の総和と一致するとき、A[v]を1増やすにはA[c]もどこか1増やさなければならない。
このことから、以下の値を考える。
f(v,i) := A[v]を1増やしたいとき、vがf(v,i)以上だと、子頂点も合わせi増やす必要があるような最小値
A[v]がA[c]の総和より大きい場合、f(c,i)を参照して増加量を抑えられる子頂点に1追加することになる。

f(v,i)をもとに、各頂点のSubTreeで条件を満たすための追加量を、DFSで順次葉から計算していけばよい。

int T,N;
ll A[5050];
int P[5050];
vector<int> C[5050];
ll ret;
vector<ll> hoge(int cur) {
	vector<ll> V={0};
	ll sum=0;
	int i;
	FORR(e,C[cur]) {
		auto W=hoge(e);
		while(V.size()<W.size()+1) V.push_back(0);
		FOR(i,W.size()) V[i+1]=min(1LL<<60,V[i+1]+W[i]);
		sum+=A[e];
	}
	if(C[cur].empty()) V={1LL<<60};
	
	
	if(A[cur]<sum) {
		V[0]=sum-A[cur];
	}
	else {
		ll d=A[cur]-sum;
		FOR(i,V.size()) {
			ll mi=min(V[i],d);
			ret+=mi*i;
			d-=mi;
			V[i]-=mi;
		}
		assert(d==0);
		
	}
	return V;
	
	
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>T;
	while(T--) {
		cin>>N;
		FOR(i,N) {
			cin>>A[i];
			C[i].clear();
		}
		for(i=1;i<N;i++) {
			cin>>P[i];
			C[P[i]-1].push_back(i);
		}
		ret=0;
		hoge(0);
		cout<<ret<<endl;
	}
}

まとめ

本番結構苦戦してるな…。