경로 가중치 합 hld 구현 코드

특징 : 맞는지 모름 앳코더에서 Proof by AC 성공

 

#include<bits/stdc++.h>
using namespace std;
#define fastio cin.tie(0)->sync_with_stdio(0)
#define X first
#define Y second
typedef pair<int,int> pii;
const int N=1e5+1;
struct Segment {
	int a[N];
	int tree[4*N],lazy[4*N];
	void init(int l,int r,int i) {
		if(l==r) { tree[i]=a[l]; return; }
		int m=l+r>>1;
		init(l,m,i<<1); init(m+1,r,i<<1|1);
		tree[i]=tree[i<<1]+tree[i<<1|1];
	}
	void prop(int l,int r,int i) {
		if(!lazy[i]) return;
		tree[i]+=lazy[i]*(r-l+1);
		if(l<r){
			lazy[i<<1]+=lazy[i];
			lazy[i<<1|1]+=lazy[i];
		}
		lazy[i]=0;
	}
	void update(int l,int r,int nl,int nr,int i,int val) {
		prop(nl,nr,i);
		if(r<nl||nr<l) return;
		if(l<=nl&&nr<=r){
			lazy[i]=val;
			prop(nl,nr,i);
			return;
		}
		int nm=nl+nr>>1;
		update(l,r,nl,nm,i<<1,val); update(l,r,nm+1,nr,i<<1|1,val);
		tree[i]=tree[i<<1]+tree[i<<1|1];
	}
	int query(int l,int r,int nl,int nr,int i) {
		prop(nl,nr,i);
		if(r<nl||nr<l) return 0;
		if(l<=nl&&nr<=r) return tree[i];
		int nm=nl+nr>>1;
		return query(l,r,nl,nm,i<<1)+query(l,r,nm+1,nr,i<<1|1);
	}
} tree;
struct HLD {
	int n, cost[N]; pair<pii,int> inp[N];
	int sz[N],dep[N],par[N],top[N],in[N],out[N];
	vector<pii> ad[N];
	vector<int> adj[N];
	void dfs(int cur,int prv) {
		for(auto nxt : ad[cur])
			if(nxt.X^prv){
				adj[cur].push_back(nxt.X);
				cost[nxt.X]=nxt.Y;
				dfs(nxt.X,cur);
			}
	}
	void dfs1(int cur) {
		sz[cur]=1;
		for(auto& nxt : adj[cur]){
			dep[nxt]=dep[cur]+1, par[nxt]=cur;
			dfs1(nxt); sz[cur]+=sz[nxt];
			if(sz[nxt]>sz[adj[cur][0]]) swap(nxt,adj[cur][0]);
		}
	}
	int tmp;
	void dfs2(int cur) {
		in[cur]=++tmp;
		for(auto nxt : adj[cur]){
			top[nxt]=(nxt==adj[cur][0]?top[cur]:nxt);
			dfs2(nxt);
		}
		out[cur]=tmp;
	}
	void build() {
		cin >> n;
		for(int i=1;i<n;i++){
			int u,v,w; cin >> u >> v >> w;
			inp[i]={{u,v},w};
			ad[u].push_back({v,w});
			ad[v].push_back({u,w});
		}
		top[1]=1;
		dfs(1,0), dfs1(1), dfs2(1);
		for(int i=1;i<=n;i++) tree.a[in[i]]=cost[i];
		tree.init(1,n,1);
	}
} hld;
void update(int i,int val) {
	int u=hld.inp[i].X.X, v=hld.inp[i].X.Y, w=hld.inp[i].Y;
	if(hld.dep[u]<hld.dep[v]) swap(u,v);
	tree.update(hld.in[u],hld.in[u],1,hld.n,1,val-w);
	hld.inp[i].Y=val;
}
int query(int a,int b) {
	int ret=0;
	while(hld.top[a]!=hld.top[b]){
		if(hld.dep[hld.top[a]]<hld.dep[hld.top[b]]) swap(a,b);
		int st=hld.top[a];
		ret+=tree.query(hld.in[st],hld.in[a],1,hld.n,1);
		a=hld.par[st];
	}
	if(a==b) return ret;
	if(hld.dep[a]>hld.dep[b]) swap(a,b);
	int nxt=0;
	for(auto i : hld.adj[a]) if(hld.top[i]==hld.top[a]) nxt=i;
	ret+=tree.query(hld.in[nxt],hld.in[b],1,hld.n,1);
	return ret;
}
int main()
{
	fastio;
	hld.build();
	int q; cin >> q;
	while(q--){
		int t,a,b; cin >> t >> a >> b;
		if(t==1) update(a,b);
		else cout << query(a,b) << '\n';
	}
}