樹鏈剖分求邊權 [USACO11DEC]Grass Planting G（洛谷 P3038）

[USACO11DEC]Grass Planting G

題目大意：
給你一顆n個結點的樹，m個操作，操作爲詢問一條路徑上的邊權和或者對一條路徑上的每條邊的權值都加1；

---

樹鏈剖分的模板是對樹上結點操作，但是這是對樹上的邊操作；

邊轉化爲點，只要兩點路徑上的每個點（除了兩點的LCA）都加1就行；

詢問和修改都是如此；

樹鏈剖分也可以求LCA，當x，y跳到同一條鏈時，深度小的就是LCA；

代碼：

#include<bits/stdc++.h>
#define LL long long
#define pa pair<int,int>
#define ls k<<1
#define rs k<<1|1
#define inf 0x3f3f3f3f
using namespace std;
const int N=1000100;
const int M=50100;
const LL mod=10007;
int n,m,head[N],cnt;
int size[N],fa[N],son[N],deep[N];
int top[N],id[N],tot;
struct node{
	int to,nex;
}edge[N*2];
struct Node{
	int l,r,lz,sum;
}tr[N*4];
void add(int p,int q){edge[cnt].to=q,edge[cnt].nex=head[p],head[p]=cnt++;}
////
void dfs1(int p,int fat){
	size[p]=1,fa[p]=fat;
	int mx=-1;
	for(int i=head[p];~i;i=edge[i].nex){
		int q=edge[i].to;
		if(q!=fat){
			deep[q]=deep[p]+1;
			dfs1(q,p);
			size[p]+=size[q];
			if(size[q]>mx) son[p]=q,mx=size[q];
		}
	}
}
void dfs2(int p,int topp){
	id[p]=++tot,top[p]=topp;
	if(!son[p]) return;
	dfs2(son[p],topp);
	for(int i=head[p];~i;i=edge[i].nex){
		int q=edge[i].to;
		if(q==fa[p]||q==son[p]) continue;
		dfs2(q,q);
	}
}
////
void pp(int k){tr[k].sum=tr[ls].sum+tr[rs].sum;}
void pd(int k){
	if(tr[k].lz){
		tr[ls].lz+=tr[k].lz,tr[rs].lz+=tr[k].lz;
		tr[ls].sum+=(tr[ls].r-tr[ls].l+1)*tr[k].lz;
		tr[rs].sum+=(tr[rs].r-tr[rs].l+1)*tr[k].lz;
	}
	tr[k].lz=0;
}
void build(int l,int r,int k){
	tr[k].l=l,tr[k].r=r;
	if(l==r) return;
	int d=(l+r)>>1;
	build(l,d,ls);
	build(d+1,r,rs);
}
void update(int l,int r,int w,int k){
	if(tr[k].l>=l&&tr[k].r<=r){
		tr[k].lz+=w;
		tr[k].sum+=(tr[k].r-tr[k].l+1)*w;
		return;
	}
	pd(k);
	int d=(tr[k].l+tr[k].r)>>1;
	if(l<=d) update(l,r,w,ls);
	if(r>d) update(l,r,w,rs);
	pp(k);
}
int query(int l,int r,int k){
	int ans=0;
	if(tr[k].l>=l&&tr[k].r<=r) return tr[k].sum;
	pd(k);
	int d=(tr[k].l+tr[k].r)>>1;
	if(l<=d) ans+=query(l,r,ls);
	if(r>d) ans+=query(l,r,rs);
	return ans;
}
////
void update_chain(int x,int y,int z){
	while(top[x]!=top[y]){
		if(deep[top[x]]<deep[top[y]]) swap(x,y);
		update(id[top[x]],id[x],z,1);
		x=fa[top[x]];
	}
	if(deep[x]<deep[y]) swap(x,y);
	update(id[y]+1,id[x],z,1);
}
int query_chain(int x,int y){
	int ans=0;
	while(top[x]!=top[y]){
		if(deep[top[x]]<deep[top[y]]) swap(x,y);
		ans+=query(id[top[x]],id[x],1);
		x=fa[top[x]];
	}
	if(deep[x]<deep[y]) swap(x,y);
	ans+=query(id[y]+1,id[x],1);
	return ans;
}
////
int main(){
	memset(head,-1,sizeof(head));
	scanf("%d%d",&n,&m);
	for(int i=1;i<n;i++){
		int x,y;scanf("%d%d",&x,&y);
		add(x,y),add(y,x);
	}
	dfs1(1,0),dfs2(1,1),build(1,n,1);
	while(m--){
		char c;int a,b;
		cin>>c>>a>>b;
		if(c=='P') update_chain(a,b,1);
		else printf("%d\n",query_chain(a,b));
	}
	return 0;
}