Codeforces #398 Div2 C. Garland

普段のDiv2より難しめ…とはいえグダグダすぎました。
http://codeforces.com/contest/767/problem/C

問題

N頂点の木を成すグラフが与えられる。
各頂点には整数が振られている。

この木の辺を2か所切断し、木を3つに分割したい。
その際、分割後の各木に含まれる頂点の値の総和を等しくしたい。
そのような分割が存在するなら一例を示せ。

解法

まず全体の値の総和を求めよう。
それが3の倍数でなければそもそも3等分は無理である。

木を根付き木とみなし、DFSで各頂点におけるSubtreeの頂点の値の総和を求めよう。
条件を満たす分割ができるのは、以下のいずれかの場合である。

• 互いに共有部分を持たないSubtreeで、Subtree内の頂点の値の総和が全体の1/3となるものが2つ以上ある。
• Subtree内の頂点の値の総和が全体の2/3であり、かつSubtree内の頂点で、さらのその頂点のSubtree内の頂点の値の総和が全体の1/3となるものがある。(ただし根は除く)

各頂点に対し、Subtree内の頂点で、さらにそのSubtree内の頂点の値の総和が全体の1/3となるものがある場合、その頂点番号を管理していくと上記両方の条件を検出できる。

int N;
vector<int> E[1010101];
int P[1010101];
ll T[1010101];

ll S;
int dfs(int cur,int pre) {
	vector<int> PP;
	P[cur]=-1;
	FORR(e,E[cur]) if(e!=pre) {
		int x=dfs(e,cur);
		if(x>=0 && PP.size()<=1) PP.push_back(x);
		T[cur]+=T[e];
	}
	
	if(PP.size()>=2) {
		_P("%d %d\n",PP[0]+1,PP[1]+1);
		exit(0);
	}
	if(PP.size()>=1 && (pre>=0 && T[cur]==S*2)) {
		_P("%d %d\n",PP[0]+1,cur+1);
		exit(0);
	}
	if(T[cur]==S) return cur;
	if(PP.size()) return PP[0];
	return -1;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	int root=0;
	FOR(i,N) {
		cin>>x>>T[i];
		x--;
		if(x>=0) E[x].push_back(i);
		else root=i;
		S+=T[i];
	}
	if(S%3) return _P("-1\n");
	S/=3;
	
	auto p=dfs(root,-1);
	_P("-1\n");
}

まとめ

総和が0のとき、根頂点とその(存在しない)親頂点を切り離してしまった…。