树链刨分模板题

题目链接

两个dfs预先处理出我们需要用到的七个信息。第一个dfs我们需要预处理出四个信息，分别是，f（结点的父亲），son（节点的重儿子），dep（节点的深度），size（节点的大小），第二个dfs处理三个信息，分别是，top（此节点所在链的根节点），id（节点在线段树中的编号），val（节点所在线段树中的点权值）。对于更新u，v这条链的值，如果u，v不在同一条链上，首先我们找到  u，v 中dep【top】大的那个节点，更新top到这个节点的值，根据第二个dfs我们可以知道，一条链肯定位于线段树的一个连续区间内。所以我们直接修改 top【id【u】】到id【u】的值，然u跳到top的父亲结点，如此往复，最后u，v会在同一条链上，就直接修改那条链的信息。根据第二个dfs我们也可以知道，以一个节点为根节点子树必定会在线段树的一段连续的区间内。如果要修改以u节点为根节点的子树的值，那么直接修改 id【u】到id【u】+size【u】-1的值即可，求和同理。

#include<bits/stdc++.h>
using namespace std;
#define ll long long
const int maxn=2e5+10; 
int a[maxn],size[maxn],f[maxn],son[maxn],dep[maxn];
int top[maxn],id[maxn],val[maxn],cnt;
int tr[maxn*4],lz[maxn*4],mod,n;
vector<int>G[maxn];
void dfs1(int u,int fa,int d)
{
	f[u]=fa;
	size[u]=1;
	dep[u]=d;
	for(int i=0;i<G[u].size();i++)
	{
		int v=G[u][i];
		if(v==f[u])continue;
		dfs1(v,u,d+1);
		size[u] += size[v];
		if(!son[u]||size[v]>size[son[u]])son[u]=v;
	}
}
void dfs2(int u,int rt)
{
	top[u]=rt; 
	id[u]=++cnt; 
	val[id[u]]=a[u];
	if(!son[u])return ;
	dfs2(son[u],rt);
	for(int i=0;i<G[u].size();i++)
	{
		int v=G[u][i];
		if(v==son[u]||v==f[u])continue;
		dfs2(v,v);
	}
} 

#define ls o<<1
#define rs o<<1|1
#define mid (l+r)/2

void build(int o,int l,int r)
{
	if(l==r)
	{
		tr[o]=val[l]%mod;
		return ;
	}
	build(ls,l,mid);
	build(rs,mid+1,r);
	tr[o]=(tr[ls]+tr[rs])%mod;
}
void down(int o,int l,int r)
{
	if(lz[o]==0)return ;
	lz[ls] = (lz[o]+lz[ls])%mod;
	lz[rs] = (lz[o]+lz[rs])%mod;
	
	tr[ls]=	(tr[ls]%mod + 1ll*(mid-l+1)*lz[o] %mod)%mod;
	tr[rs]= (tr[rs]%mod + 1ll*(r-mid)*lz[o] %mod)%mod;
	lz[o]=0;
}
void up(int o,int l,int r,int L,int R,int v)
{
	if(l>=L&&r<=R)
	{
		lz[o]= (lz[o]+ v%mod)%mod;
		tr[o] = (tr[o]%mod +1ll*(r-l+1)*v%mod)%mod ;
		return ;
	}
	down(o,l,r);
	if(L<=mid)up(ls,l,mid,L,R,v);
	if(R>mid)up(rs,mid+1,r,L,R,v);
	tr[o]=(tr[ls]+tr[rs])%mod;
}

long long qu(int o,int l,int r,int L,int R)
{
	if(l>=L&&r<=R)
	{
		return tr[o]%mod;	
	}
	down(o,l,r);
	long long ans=0;
	if(mid>=L)ans= (ans + qu(ls,l,mid,L,R)%mod)%mod;
	if(R>mid)ans=(ans+ qu(rs,mid+1,r,L,R)%mod)%mod;
	return ans%mod;
}
void up_path(int u,int v,int c)
{
	while(top[u]!=top[v])
	{
		if(dep[top[u]]<dep[top[v]])swap(u,v);
		up(1,1,cnt,id[top[u]],id[u],c);
		u=f[top[u]];
	}
	if(id[u]>id[v])swap(u,v);
	up(1,1,cnt,id[u],id[v],c);
}

long long qu_path(int u,int v)
{
	long long ans=0;
	while(top[u]!=top[v])
	{
		if(dep[top[u]]<dep[top[v]])swap(u,v);
		ans = (ans+qu(1,1,cnt,id[top[u]],id[u])%mod)%mod;
		u=f[top[u]];
	}
	if(id[u]>id[v])swap(u,v);
	ans+=qu(1,1,cnt,id[u],id[v])%mod;
	return ans%mod;
}
int main()
{
	int n,qq,root;
	scanf("%d %d %d %d",&n,&qq,&root,&mod);
	for(int i=1;i<=n;i++)scanf("%d",a+i);
	for(int i=1;i<=n-1;i++)
	{
		int u,v;
		scanf("%d %d",&u,&v);
		G[u].push_back(v);
		G[v].push_back(u);
	}
	dfs1(root,0,1);
	dfs2(root,root);
	build(1,1,n);
	while(qq--)
	{
		int opt;
		scanf("%d",&opt);
		if(opt==1)
		{
			int x,y,v;
			scanf("%d %d %d",&x,&y,&v);
			up_path(x,y,v%mod);
		}
		else if(opt==2)
		{
			int x,y;
			scanf("%d %d",&x,&y);
			long long ans=qu_path(x,y);
			printf("%lld\n",ans); 
		}
		else if(opt==3)
		{
			int x,z;
			scanf("%d %d",&x,&z);
			up(1,1,cnt,id[x],id[x]+size[x]-1,z%mod);
		}
		else 
		{
			int x;
			scanf("%d",&x);
			printf("%d\n",qu(1,1,cnt,id[x],id[x]+size[x]-1));
		}
	}
	return 0;
 }