AtCoder ABC #133 : F - Colorful Tree

もうちょいサクサク行きたかった。
https://atcoder.jp/contests/abc133/tasks/abc133_f

問題

N頂点の木を成す無向グラフが与えられる。
各辺には色と長さが設定されている。

以下のクエリに順次答えよ。
なお、クエリによる辺の長さの変化は、以降のクエリのときにはもとに戻る。

• 色Xの辺の長さがすべてYになったとき、U,V間の距離を求めよ。

解法

木における頂点間の距離の求め方の定番として、
D(v) := 根から頂点vへの距離
としたとき、 D(U)+D(V)-2*D(LCA(U,V)) を取るというものがある。
これを応用しよう。
U-V間にある色Xの辺の数NXと、その辺の長さの総和をSLとする。
すると解は、元の長さから色Xの辺による分を引き、長さを変えた色Xの辺の長さを足すとD(U)+D(V)-2*D(LCA(U,V))-SL+NX*Yとなる。
あとはNXとLを求めればよい。

さて、頂点間における特定の色の登場頻度と距離の総和だが、

• C(c,v) := 根から頂点vに至る経路における色cの辺の数
• L(c,v) := 根から頂点vに至る経路における色cの辺の長さの総和

とすると、距離と同じ感じで

• NX=C(X,U)+C(X,V)-2*C(X,LCA(U,V))
• SL=L(X,U)+L(X,V)-2*L(X,LCA(U,V))

で求められる。

C(c,v)とL(c,v)を全部埋めようとすると、TLEやMLEするが、各クエリにおいて、U,V,LCA(U,V)の3頂点における値さえわかればよい。
よって先にクエリを先読みし、各頂点vで、C(c,v)とL(c,v)を求めておくべき色を覚えておこう。

あとはDFSで根頂点から各頂点に至る辺の各色の数・距離の総和を更新しながら、訪問した頂点vで必要なC(c,v)とL(c,v)だけ覚えておく。
これで必要な値はすべて求まるので、各クエリに答えていく。

（なお、DFS中にクエリに必要な加減算をしてしまった方が実装が単純）

int N,Q;
vector<pair<int,int>> E[202020];
int A[101010],B[101010],C[101010],D[101010];
map<int,pair<int,ll>> memo[101010];
int num[101010];
ll dist[101010];
ll TD[101010];

int X[101010],Y[101010],U[101010],V[101010];
int P[21][200005],DD[200005];

void dfs(int cur) {
	ITR(it,E[cur]) if(it->first!=P[0][cur]) {
		TD[it->first]=TD[cur]+D[it->second];
		DD[it->first]=DD[cur]+1, P[0][it->first]=cur, dfs(it->first);
	}
}

int lca(int a,int b) {
	int ret=0,i,aa=a,bb=b;
	if(DD[aa]>DD[bb]) swap(aa,bb);
	for(i=19;i>=0;i--) if(DD[bb]-DD[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 dfs2(int cur,int pre) {
	FORR(m,memo[cur]) {
		m.second={num[m.first],dist[m.first]};
	}
	FORR(e,E[cur]) if(e.first!=pre) {
		int v=e.second;
		num[C[v]]++;
		dist[C[v]]+=D[v];
		dfs2(e.first,cur);
		num[C[v]]--;
		dist[C[v]]-=D[v];
		
	}
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>Q;
	FOR(i,N-1) {
		cin>>A[i]>>B[i]>>C[i]>>D[i];
		A[i]--;
		B[i]--;
		E[A[i]].push_back({B[i],i});
		E[B[i]].push_back({A[i],i});
	}
	dfs(0);
	FOR(i,19) FOR(x,N) P[i+1][x]=P[i][P[i][x]];
	FOR(i,Q) {
		cin>>X[i]>>Y[i]>>U[i]>>V[i];
		U[i]--;
		V[i]--;
		memo[U[i]][X[i]]={0,0};
		memo[V[i]][X[i]]={0,0};
		x=lca(U[i],V[i]);
		memo[x][X[i]]={0,0};
	}
	dfs2(0,-1);
	FOR(i,Q) {
		x=lca(U[i],V[i]);
		ll cd=(TD[U[i]]-memo[U[i]][X[i]].second)+(TD[V[i]]-memo[V[i]][X[i]].second)-2*(TD[x]-memo[x][X[i]].second);
		int tnum=memo[U[i]][X[i]].first+memo[V[i]][X[i]].first-2*memo[x][X[i]].first;
		cout<<cd+1LL*tnum*Y[i]<<endl;
	}
	
	
}

まとめ

• DFS順であれば経由した辺の数・距離がわかる
• クエリで必要な情報はそのうち一部

ってのが重要だね。

なんかあわてて解いたから変数名がメチャクチャだ。