树链剖分-树剖换根

Description

这是一道模板题。
给定一棵 n 个节点的树，初始时该树的根为 1 号节点，每个节点有一个给定的权值。下面依次进行 m 个操作，操作分为如下五种类型：
换根：将一个指定的节点设置为树的新根。
修改路径权值：给定两个节点，将这两个节点间路径上的所有节点权值（含这两个节点）增加一个给定的值。
修改子树权值：给定一个节点，将以该节点为根的子树内的所有节点权值增加一个给定的值。
询问路径：询问某条路径上节点的权值和。
询问子树：询问某个子树内节点的权值和。

Input

第一行为一个整数 nnn，表示节点的个数。
第二行 nnn 个整数表示第 iii 个节点的初始权值 ai​​ 。
第三行 n−1 个整数，表示 i+1号节点的父节点编号 fi+1 (1⩽fi+1⩽n)。
第四行一个整数 m，表示操作个数。
接下来 m 行，每行第一个整数表示操作类型编号：(1⩽u,v⩽n)
若类型为 1，则接下来一个整数 u，表示新根的编号。
若类型为 2，则接下来三个整数 u,v,k，分别表示路径两端的节点编号以及增加的权值。
若类型为 3，则接下来两个整数 u,k，分别表示子树根节点编号以及增加的权值。
若类型为 4，则接下来两个整数 u,v，表示路径两端的节点编号。
若类型为 5，则接下来一个整数 u，表示子树根节点编号。

Output

对于每一个类型为 4 或 5 的操作，输出一行一个整数表示答案。

---

思路

• 重构两遍40分，后发现getson函数写错，正确：if(fa[top[u]]==v) return top[u];错误：if(fa[u]==v) return u;
• 树剖换根相关： root(当前根)，x(询问)，lca最近公共祖先，以询问总和sum为例：
• 1. root==x：sum=整棵树之和；
• 2. root在x的子树中(lca(root,x)==x)：sum=整棵树之和-以(x~root路径上靠近x的点)为根的子树；
• 3. root不在x的子树中(lca(root,x)!=x)：sum=以1为根(原树)时x的子树；
• 以1为根时x的子树在线段树中为连续的一段区间(num[x],num[x]+siz[x]-1)；

---

代码

#include <iostream>
#include <cstdio>
#define int long long
using namespace std;
const int maxn=100003;
int n,m,a[maxn],head[maxn],cnt,root,sum[maxn<<2],tag[maxn<<2],len[maxn<<2];
int dep[maxn],fa[maxn],top[maxn],num[maxn],fnum[maxn],son[maxn],siz[maxn];
struct node{int next,to;}e[maxn];
void addedge(int x,int y){e[++cnt]=(node){head[x],y},head[x]=cnt;}
void dfs_1(int x){
	siz[x]=1,dep[x]=dep[fa[x]]+1;
	for(int i=head[x],v;v=e[i].to,i;i=e[i].next){
		dfs_1(v),siz[x]+=siz[v];
		if(!son[x]||siz[v]>siz[son[x]]) son[x]=v;
	}
}
void dfs_2(int x,int topp){
	top[x]=topp,num[x]=++cnt,fnum[cnt]=x;
	if(son[x]) dfs_2(son[x],topp);
	for(int i=head[x];i;i=e[i].next) if(e[i].to!=son[x]) dfs_2(e[i].to,e[i].to);
}
void gettag(int x,int d){tag[x]+=d,sum[x]+=d*len[x];}
void pushdown(int x){if(tag[x]) gettag(x<<1,tag[x]),gettag(x<<1|1,tag[x]),tag[x]=0;}
void pushup(int x){sum[x]=sum[x<<1]+sum[x<<1|1];}
void build(int x,int lef,int rig){
	len[x]=rig-lef+1;
	if(lef==rig) return sum[x]=a[fnum[lef]],void();
	int mid=lef+rig>>1;
	build(x<<1,lef,mid),build(x<<1|1,mid+1,rig);
	pushup(x);
}
void add(int x,int xl,int xr,int lef,int rig,int d){
	if(xr<lef||xl>rig) return;
	if(lef<=xl&&rig>=xr) return gettag(x,d),void();
	int mid=xl+xr>>1;
	pushdown(x);
	add(x<<1,xl,mid,lef,rig,d),add(x<<1|1,mid+1,xr,lef,rig,d);
	pushup(x);
}
int ask(int x,int xl,int xr,int lef,int rig){
	if(xr<lef||xl>rig) return 0;
	if(lef<=xl&&rig>=xr) return sum[x];
	pushdown(x);
	int mid=xl+xr>>1;
	return ask(x<<1,xl,mid,lef,rig)+ask(x<<1|1,mid+1,xr,lef,rig);
}
int getlca(int u,int v){
	while(top[u]!=top[v]){
		if(dep[top[u]]<dep[top[v]]) swap(u,v);
		u=fa[top[u]];
	}return dep[u]<dep[v]?u:v;
}
int getson(int u,int v){
	while(top[u]!=top[v]){
		if(dep[top[u]]<dep[top[v]]) swap(u,v);
		if(fa[top[u]]==v) return top[u];
		u=fa[top[u]];
	}return dep[u]<dep[v]?son[u]:son[v];
}
void modify_uv(int u,int v,int k){
	while(top[u]!=top[v]){
		if(dep[top[u]]<dep[top[v]]) swap(u,v);
		add(1,1,n,num[top[u]],num[u],k);
		u=fa[top[u]];
	}
	if(dep[u]>dep[v]) swap(u,v);
	add(1,1,n,num[u],num[v],k);
}
void modify_u(int u,int k){
	int lca=getlca(root,u),son;
	if(u==root) return add(1,1,n,1,n,k);
	if(u!=lca) return add(1,1,n,num[u],num[u]+siz[u]-1,k);
	else add(1,1,n,1,n,k),son=getson(u,root),add(1,1,n,num[son],num[son]+siz[son]-1,-k);
}
int query_uv(int u,int v){
	int ans=0;
	while(top[u]!=top[v]){
		if(dep[top[u]]<dep[top[v]]) swap(u,v);
		ans+=ask(1,1,n,num[top[u]],num[u]);
		u=fa[top[u]];
	}if(dep[u]>dep[v]) swap(u,v);
	return ans+ask(1,1,n,num[u],num[v]);
}
int query_u(int u){
	int lca=getlca(root,u),son;
	if(u==root) return ask(1,1,n,1,n);
	if(u!=lca) return ask(1,1,n,num[u],num[u]+siz[u]-1);
	else return son=getson(u,root),ask(1,1,n,1,n)-ask(1,1,n,num[son],num[son]+siz[son]-1);
}
signed main()
{
	scanf("%lld",&n);
	for(int i=1;i<=n;++i) scanf("%lld",&a[i]);
	for(int i=2,x;i<=n;++i) scanf("%lld",&fa[i]),addedge(fa[i],i);
	root=1,dfs_1(1),cnt=0,dfs_2(1,1),build(1,1,n);
	scanf("%lld",&m);
	while(m--){
		int opt,u,v,k;scanf("%lld%lld",&opt,&u);
		if(opt==1) root=u;
		else if(opt==2) scanf("%lld%lld",&v,&k),modify_uv(u,v,k);
		else if(opt==3) scanf("%lld",&k),modify_u(u,k);
		else if(opt==4) scanf("%lld",&v),cout<<query_uv(u,v)<<'\n';
		else cout<<query_u(u)<<'\n';
	}
	return 0;
}