kmjp's blog

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

Codeforces #396 Div2 E. Mahmoud and a xor trip

これも割と定番?
http://codeforces.com/contest/766/problem/E

問題

木を成す無向グラフが与えられる。
各頂点には整数が振られている。
ここで2つの端点を選びパスを作ることを考える(両端点は同じ頂点でもよい)。

パス上の頂点の値のxorをとったものを考える。
全パスにおけるその値の総和を求めよ。

解法

xorを取る問題なのでbit単位で総和を取ればよい。
各頂点において、その頂点を通りかつその頂点のSubtree内に閉じたパスに対する値を数え上げよう。
その頂点vの子頂点cに対する各SubTreeにおいて、vに至るまで各bitが1になるものがいくつあるかを数え上げよう。
あるSubTreeでvまでのxorを取るとp bit目が1になる頂点がa個あり、別のSubTreeでcまでのxorを取るとp bit目が0になる頂点がb個あるなら、a*b*2^p分だけ解に寄与する。
あとは各SubTree内で各bitが0,1となるものがいくつあるかを数え上げていけばよい。

int N;
int A[101010];
vector<int> E[101010];
int C[101010];
ll num[101010][21];
ll ret;

void dfs(int cur,int pre) {
	int i;
	C[cur]=1;
	ret += A[cur];
	FOR(i,20) if(A[cur]&(1<<i)) num[cur][i]++;
	FORR(e,E[cur]) if(e!=pre) {
		dfs(e,cur);
		FOR(i,20) {
			ret += (num[cur][i]*(C[e]-num[e][i])+(C[cur]-num[cur][i])*num[e][i])<<i;
			if(A[cur]&(1<<i)) {
				num[cur][i] += C[e]-num[e][i];
			}
			else {
				num[cur][i] += num[e][i];
			}
		}
		C[cur] += C[e];
	}
	
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	FOR(i,N) cin>>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<<ret<<endl;
}

まとめ

最初2回DFSして全方向の0/1登場回数を数えようとしたけど、その必要はなかったようだ。