kmjp's blog

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

AtCoder ABC #294 : G - Distance Queries on a Tree

こちらは典型かな…。
https://atcoder.jp/contests/abc294/tasks/abc294_g

問題

木を成す無向グラフが与えられる。
各辺には重さが設定されている。

以下のクエリに順次答えよ。

  • 指定された辺の重みを指定した値に変える。
  • 指定された2点(u,v)間の最短パス上の辺の重みの総和を答える。

解法

f(v) := 根頂点からvまでの最短パス上の辺の重みの総和
とすると、後者の解はf(u)+f(v)-2*f(LCA(u,v))である。
あとはf(v)を高速に求められれば良い。

先に頂点をDFS順に並べておくと、点Parent(a)-a間に重さwの辺があった場合、aのsubtreeの各点vにおいて、f(v)にwが加算される。
これは区間加算に相当するので、BITなりSegTreeで計算できる。

int N,Q;
int U[202020],V[202020],W[202020];
int L[202020],R[202020];
int id;
vector<int> E[202020];

template<class V, int ME> class BIT {
public:
	V bit[1<<ME];
	V operator()(int e) {if(e<0) return 0;V s=0;e++;while(e) s+=bit[e-1],e-=e&-e; return s;}
	void add(int e,V v) { e++; while(e<=1<<ME) bit[e-1]+=v,e+=e&-e;}
};
BIT<ll,20> bt;

int P[21][200005],D[200005];

void dfs(int cur) {
	L[cur]=id++;
	FORR(e,E[cur]) if(e!=P[0][cur]) D[e]=D[cur]+1, P[0][e]=cur, dfs(e);
	R[cur]=id;
}
int getpar(int cur,int up) {
	int i;
	FOR(i,20) if(up&(1<<i)) cur=P[i][cur];
	return cur;
}

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; string s;
	
	cin>>N;
	FOR(i,N-1) {
		cin>>U[i]>>V[i]>>W[i];
		U[i]--;
		V[i]--;
		E[U[i]].push_back(V[i]);
		E[V[i]].push_back(U[i]);
	}
	dfs(0);
	FOR(i,19) FOR(x,N) P[i+1][x]=P[i][P[i][x]];
	
	FOR(i,N-1) {
		if(L[U[i]]>L[V[i]]) swap(U[i],V[i]);
		bt.add(L[V[i]],W[i]);
		bt.add(R[V[i]],-W[i]);
	}
	cin>>Q;
	while(Q--) {
		cin>>i>>x>>y;
		if(i==1) {
			x--;
			bt.add(L[V[x]],-W[x]);
			bt.add(R[V[x]],W[x]);
			W[x]=y;
			bt.add(L[V[x]],W[x]);
			bt.add(R[V[x]],-W[x]);
		}
		else {
			x--,y--;
			int lc=lca(x,y);
			ll a=bt(L[lc]);
			ll b=bt(L[x]);
			ll c=bt(L[y]);
			cout<<b+c-2*a<<endl;
		}
	}
	
	
}

まとめ

Gにしては典型寄りかも。