kmjp's blog

競技プログラミング参加記です

Codeforces #530 Div1 A. Sum in the tree

最近中難易度の問題によく取り組むので、ブログ書く方が遅れている。
https://codeforces.com/contest/1098/problem/A

問題

根付き木が与えられる。
各点には、設定値S[i]が指定されている。ただし高さが偶数の位置にある頂点ではS[i]は不定である。

この状態で、各頂点に非負整数A[i]を振ることを考える。
各頂点vに対し、root→vのパス上のA[i]の総和を取ると、S[i]となるようにしたい。
また、そのうちでA[i]の総和を最小にせよ。

解法

本番、「偶数の高さの~」を読み飛ばしてしまったが、それでも解くことはできる。
まず、祖先-子孫の関係にある頂点間で、祖先の方が大きなS[i]を持つことは許容されない。

根に近い方のA[i]を増やす方が、A[i]の総和は最終的に小さくなることが期待できる。
この条件を元に、DFSで未確定のS[i]を埋めていこう。
根に近い側のA[i]を増やすためには、S[i]も根に近い方が大きくなればよい。

Sがすべて定まればAも確定する。

int N;
int P[101010];
vector<int> E[101010];
ll S[101010];
ll MS[101010];
ll A[101010];

void solve() {
	int i,j,k,l,r,x,y; string s;
	cin>>N;
	for(i=1;i<N;i++) {
		cin>>P[i];
		P[i]--;
		E[P[i]].push_back(i);
	}
	FOR(i,N) cin>>S[i];
	
	for(i=N-1;i>=0;i--) {
		MS[i]=-1;
		FORR(e,E[i]) if(MS[e]!=-1) {
			if(MS[i]==-1) MS[i]=MS[e];
			else MS[i]=min(MS[i],MS[e]);
		}
		
		if(MS[i]==-1) {
			MS[i]=S[i];
		}
		else {
			if(S[i]==-1) {
				S[i]=MS[i];
			}
			else {
				if(MS[i]<S[i]) return _P("-1\n");
				MS[i]=S[i];
			}
		}
	}
	
	if(S[0]==-1) return _P("0\n");
	
	ll sum=0;
	FOR(i,N) {
		if(S[i]>-1) {
			if(i==0) {
				A[i]=S[i];
			}
			else {
				A[i]=S[i]-S[P[i]];
			}
		}
		else {
			S[i]=S[P[i]];
		}
		sum+=A[i];
	}
	cout<<sum<<endl;
	
}

まとめ

いや、Aでこんなに書いておかしいと思ったんだよね。