Codeforces #202 Div1. B. Apple Tree

本番オーバーフローしてRuntime Errorさせてしまった問題。
http://codeforces.com/contest/348/problem/B

問題

根付き木があり、各点にはりんごの数が割り振られている。
ここで、各点を頂点とするsubtreeの重さは、各点およびsubtreeのりんごの数の総和だとする。

木がバランスしているとは、各点において各subtreeの重さが一致している状態をいう。
与えられた根付き木をバランスさせるために取り除く必要がある最少のりんごの数を答えよ。

解法

DFSで葉から順にバランスさせていく。
ある点Vが3つ子W1,W2,W3を持っている場合、その親の点PをバランスさせるためにVの重さを減らさないといけないとすると、Vのバランスを保つにはW1,W2,W3から同じ数、すなわち各点1個の計3個ずつ減らさないといけない。
さらに一般化すると、ある点の重さの減らし方は、その個の点の減らし方の最小公倍数の倍数でなければならない。

よって、各点において重さと減らし方の2値を計算してDFSすればよい。
減らし方は最小公倍数を頻繁に計算するため、64bitでもオーバーフローする場合がある。
しかしリンゴの数はそもそも10^13しかないので、減らし方がそれを超えそうなら実質減らせないと考えることができる。

int N;
ll TO;
int A[100003];
ll T[100003],M[100003];

vector<int> E[100001];

void dfs(int cur,int pre) {
	int i,v=0;
	M[cur]=1;
	T[cur]=A[cur];

	if(cur!=0 && E[cur].size()==1) return;
	
	FOR(i,E[cur].size()) {
		int tar=E[cur][i];
		if(tar==pre) continue;
		v++;
		dfs(tar,cur);
		T[cur]+=T[tar];
		M[cur] = M[cur]*M[tar]/__gcd(M[cur],M[tar]);
		if(M[cur]>T[cur]) {
			T[cur]=0;
			return;
		}
	}
	
	
	ll mo=1LL<<60;
	FOR(i,E[cur].size()) {
		int tar=E[cur][i];
		if(tar==pre) continue;
		T[tar] -= T[tar] % M[cur];
		mo=min(mo,T[tar]);
	}
	T[cur]=v*mo;
	M[cur] *= v;
	
}

void solve() {
	int f,i,j,k,x,y,z,tx,ty;
	
	cin>>N;
	FOR(i,N) cin>>A[i];
	FOR(i,N) TO+=A[i];
	
	FOR(i,N-1) {
		cin>>x>>y;
		E[x-1].push_back(y-1);
		E[y-1].push_back(x-1);
	}
	
	dfs(0,-1);
	cout << TO-T[0] << endl;
	
	return;
}

まとめ

本番、オーバーフローの可能性に気づきつつもすでにLock後だったのよね…。