树链剖分求边权 [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;
}