kmjp's blog

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

Codeforces #282 Div1 D. Birthday

Codeforces

こちらも1WAしたものの自力回答。
http://codeforces.com/contest/494/problem/D

問題

距離付無向辺をもつN頂点の根付き木が与えられる。
1番の頂点を根とする。

各記号は以下の意味を示す：

  • d(u,v): 頂点u,v間の最短距離
  • S(v): d(1,u)=d(1,v)+d(v,u)を満たす頂点uの集合
  • f(u,v): \displaystyle \sum_{x \in S(v)} d(u,x)^2 - \sum_{x \notin S(v)} d(u,x)^2

Q個のクエリU[i]、V[i]が与えられるので、f(U[i],V[i])を答えよ。

解法

最後の式を変形すると、２つ目の総和においてvを無視できるようになる。
\displaystyle f(u,v) = \sum_{x \in S(v)} d(u,x)^2 - \sum_{x \notin S(v)} d(u,x)^2 = 2 \times \sum_{x \in S(v)} d(u,x)^2 - \sum_{x} d(u,x)^2

よって、まずは後者の各頂点に対する全頂点への距離の２乗和を求めよう。
これはCodeforcesにはしばしば出てくる問題で、DFSを２回行いサブツリーの頂点数・距離の和・距離の２乗和を列挙できる。

例えば頂点xから距離dの子頂点yがあり、yのサブツリーの頂点数がN(y)、距離の和をS(y)、二乗和をT(y)とする。
N(y)==3であり、yからサブツリーの各頂点との距離をp1,p2,p3の場合を考えるとS(y)=p1+p2+p3、T(y)=p1^2+p2^2+p3^2である。
頂点xからこの3頂点を見ると、距離の和は(d+p1)+(d+p2)+(d+p1)=N(y)*d+S(y)、二乗和は(d+p1)^2+(d+p2)^2+(d+p3)^2=N(y)*d^2 + 2*d*(p1+p2+p3) + p1^2+p2^2+p3^2 = N(y)*d^2 + 2*d*S(y)+ T(y)と、具体的なp1,p2,p3が不明でもS(y)とT(y)から距離の和と二乗和を求められる。

上記処理はサブツリー方向の距離の和と二乗和しか計算しないので、２回目のDFSで親側の距離の和と二乗和を加算してやればよい。


次に前者の \displaystyle \sum_{x \in S(v)} d(u,x)^2の計算である。
S(v)の条件が一見ややこしいが、これは単にvのサブツリーを示す。
あとはクエリにおいてU[i]がV[i]のサブツリーかどうかで処理が分岐する。

  • U[i]がV[i]のサブツリー外：前処理の段階でV[i]のサブツリーにおける距離の和と二乗和が求めてある。その値から、d(U[i],V[i])の値を用いてU[i]とV[i]のサブツリーの各点における距離の和と二乗和を求めればよい。
  • U[i]がV[i]のサブツリー内：U[i]から他の全頂点への距離の２乗和から、d(U[i],V[i])の値を用いてV[i]の親側の頂点への距離の二乗和を取ればよい。
    • d(U[i],V[i])はLCAを使って計算する。
int N,Q;
vector<pair<int,ll> > E[102000];
ll mo=1000000007;
ll num[102000];
ll tot1[102000],tot2[102000];
ll tota1[102000],tota2[102000];
ll dist[102000];
int P[21][1<<20],D[1<<20];

void dfs1(int cur,int pre) {
	int i;
	P[0][cur]=pre;
	num[cur]=1;
	FOR(i,E[cur].size()) {
		int tar=E[cur][i].first;
		if(pre==tar) continue;
		D[tar]=D[cur]+1;
		dist[tar]=(dist[cur]+E[cur][i].second)%mo;
		dfs1(tar,cur);
		num[cur]+=num[tar];
		tot1[cur]+=(tot1[tar]+num[tar]*E[cur][i].second)%mo;
		tot2[cur]+=(tot2[tar]+2*E[cur][i].second%mo*tot1[tar]+num[tar]*E[cur][i].second%mo*E[cur][i].second)%mo;
	}
	tot1[cur]%=mo;
	tot2[cur]%=mo;
}

void dfs2(int cur,int pre,ll t1,ll t2) {
	int i;
	tota1[cur]=(tot1[cur]+t1)%mo;
	tota2[cur]=(tot2[cur]+t2)%mo;
	FOR(i,E[cur].size()) {
		int tar=E[cur][i].first;
		if(pre==tar) continue;
		ll t1n=tota1[cur]+(mo-(tot1[tar]+num[tar]*E[cur][i].second)%mo)+(N-num[tar])*E[cur][i].second;
		ll t2n=tota2[cur]+(mo-(tot2[tar]+2*E[cur][i].second%mo*tot1[tar]+num[tar]*E[cur][i].second%mo*E[cur][i].second)%mo)
			+2*E[cur][i].second%mo*(tota1[cur]+(mo-(tot1[tar]+num[tar]*E[cur][i].second)%mo))+(N-num[tar])*E[cur][i].second%mo*E[cur][i].second;
		dfs2(tar,cur,t1n%mo,t2n%mo);
	}
}

int lca(int a,int b) {
	int ret=0,i,aa=a,bb=b;
	if(D[aa]>D[bb]) swap(aa,bb);
	for(i=19;i>=0;i--) if(D[bb]-D[aa]>=1<<i) bb=P[i][bb];
	for(i=19;i>=0;i--) if(P[i][aa]!=P[i][bb]) aa=P[i][aa], bb=P[i][bb];
	return (aa==bb)?aa:P[0][aa];               // vertex
}

void solve() {
	int i,j,k,l,r,x,y,u,v; string s;
	
	cin>>N;
	FOR(i,N-1) {
		cin>>x>>y>>j;
		E[x-1].push_back(make_pair(y-1,j));
		E[y-1].push_back(make_pair(x-1,j));
	}
	dfs1(0,-1);
	dfs2(0,-1,0,0);
	
	FOR(i,20) FOR(j,1<<20) P[i+1][j]=P[i][P[i][j]];
	
	cin>>Q;
	while(Q--) {
		cin>>u>>v;
		u--,v--;
		
		ll t1=0;
		x=lca(u,v);
		if(x==v) {
			// u is in v's subtree
			ll d=(dist[u]-dist[v])%mo;
			t1=tota2[u]-((N-num[v])*d%mo*d%mo+(tota2[v]-tot2[v])+2*d*(tota1[v]-tot1[v]))%mo;
		}
		else {
			// u is out of v's subtree
			ll d=(dist[u]+dist[v]-2*dist[x])%mo;
			t1=num[v]*d%mo*d%mo+2*tot1[v]*d+tot2[v];
		}
		ll r=(2*t1-tota2[u])%mo;
		if(r<0) r+=mo;
		cout << r << endl;
	}
	
}

まとめ

木の距離の二乗和とLCAの連携ということで、新しさはないが面倒な問題。
S(v)の定義とか無駄に問題文がややこしいなぁ。