kmjp's blog

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

Codeforces #474 E. Alternating Tree

手間取ったうえにしょうもないミスで落とした。
https://codeforces.com/contest/960/problem/E

問題

木を成す無向グラフが与えられる。
各頂点iには整数値V[i]が設定されている。

頂点U(1)とU(m)をつなぐパス上にある頂点列U(1)~U(m)を考える。
このパスのスコアは A(u_{1},u_{m}) = \sum\limits_{i=1}^{m} (-1)^{i+1} \cdot V_{u_{i}}とする。

全頂点対(始点と終点を逆にしたものは別々に数える)を両端をするパスに対し、スコアの総和を求めよ。

解法

各頂点の値が何回プラスおよびマイナスでカウントされるかを全方位DPで求めよう。
各頂点がプラス・マイナスどちらでカウントされるかは、始点からの距離の偶奇で決まる。
頂点vの隣接点のうち1つをwとする。
vから見てw側にある頂点のうち、vとの距離が偶数となる頂点の数をNe、奇数となる頂点の数をNoとすると、終点はv側の頂点の数だけあり得るのでV[v] * (N-(No+Ne))*(Ne-No)だけ答えにカウントされる。
なお、始点がvとなるケースは終点がN頂点のいずれでもよいので注意。

int N;
int D[202020];
ll V[202020];
int C[202020][2];
vector<int> E[202020];
ll ret;
ll mo=1000000007;

void dfs(int cur,int pre,int d) {
	D[cur]=d;
	C[cur][d]=1;
	
	FORR(e,E[cur]) if(e!=pre) {
		dfs(e,cur,d^1);
		C[cur][0]+=C[e][0];
		C[cur][1]+=C[e][1];
	}
	
}
void dfs2(int cur,int pre) {
	int CC[2]={C[0][0],C[0][1]};
	CC[D[cur]]--;
	
	FORR(e,E[cur]) if(e!=pre) {
		dfs2(e,cur);
		
		ll ot=N-C[e][0]-C[e][1];
		(ret+=(V[cur]*ot%mo)*(C[e][D[cur]]-C[e][D[cur]^1]+mo))%=mo;
		
		CC[0]-=C[e][0];
		CC[1]-=C[e][1];
	}
	(ret+=(V[cur]*(N-CC[0]-CC[1])%mo)*(CC[D[cur]]-CC[D[cur]^1]+mo))%=mo;
	(ret+=(V[cur]*N%mo))%=mo;
	
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	FOR(i,N) {
		cin>>V[i];
		V[i]=(V[i]+mo)%mo;
	}
	
	FOR(i,N-1) {
		cin>>x>>y;
		E[x-1].push_back(y-1);
		E[y-1].push_back(x-1);
	}
	dfs(0,0,0);
	dfs2(0,0);
	cout<<ret<<endl;
	
}

まとめ

自分も含め、予想以上に多くの人が落とした。