树链剖分基础模板

#include<bits/stdc++.h>
#define il inline
#define pb push_back
#define ms(_data,v) memset(_data,v,sizeof(_data))
#define sc(n) scanf("%d",&n)
#define SC(n,m) scanf("%d %d",&n,&m)
#define SZ(a) int((a).size())
#define rep(i,a,b) for(int i=a;i<=b;++i)
#define drep(i,a,b)	for(int i=a;i>=b;--i)
using namespace std;
typedef long long ll;
const ll inf=0x3f3f3f3f;
const double pi=acos(-1.0);
const double eps=1e-9;
const int maxn=1e5+5;
int n,m,r;
ll psum[maxn],sump[maxn];
vector<int> eg[maxn];
/* BIT */ 
void update(int p,ll x) {
	int pp=p;
	while(pp<=n) psum[pp]+=x*p,sump[pp]+=x,pp+=pp&-pp;	
}
void add(int l,int r,ll x) {
	update(l,x),update(r+1,-x);
}
ll query(int p) {
	ll sum=0,pp=p;
	while(pp) sum=sump[pp]*(p+1)-psum[pp],pp-=pp&-pp;
	return sum;
}
ll query(int l,int r) {
	return query(r)-query(l-1);
}
int dep[maxn],fa[maxn],sz[maxn],son[maxn],id[maxn];
int w[maxn],nw[maxn],bel[maxn],cnt=0;
/* TCP */ 
void dfs1(int x,int ff){
	dep[x]=dep[ff]+1;
	fa[x]=ff,sz[x]=1;
	int mx=-1;
	for(int i=0;i<SZ(eg[x]);++i){
		int to=eg[x][i];
		if(to==ff) continue;
		dfs1(to,x);
		sz[x]+=sz[to];
		if(sz[to]>mx){
			son[x]=to,mx=sz[to]; 
		} 
	}
}
void dfs2(int x,int topx){
	id[x]=++cnt;
	nw[cnt]=w[x];
	bel[x]=topx;
	if(!son[x]) return;
	dfs2(son[x],topx); //重儿子 
	for(int i=0;i<SZ(eg[x]);++i){ //轻儿子 
		int to=eg[x][i];
		if(to==fa[x] || to==son[x]) continue;
		dfs2(to,to);
	}
}
void r_update(int x,int y,ll z){ //将x~y最短距离上的点+z 
	while(bel[x]!=bel[y]){
		if(dep[bel[x]]<dep[bel[y]]) swap(x,y);
		add(id[bel[x]],id[x],z);
		x=fa[bel[x]];
	}
	if(dep[x]>dep[y]) swap(x,y);
	add(id[x],id[y],z);
} 
void p_update(int x,ll z){ //将x的子树都+z 
	add(id[x],id[x]+sz[x]-1,z);
} 
ll r_ask(int x,int y){ // 询问x~y最短距离上的点值和 
	ll res=0;
	while(bel[x]!=bel[y]){
		if(dep[bel[x]]<dep[bel[y]]) swap(x,y);
		res+=query(id[bel[x]],id[x]);
		x=fa[bel[x]];
	}
	if(dep[x]>dep[y]) swap(x,y);
	res+=query(id[x],id[y]); 
	return res;
}
ll p_ask(int x){ //询问x的子树点值和 
	return query(id[x],id[x]+sz[x]-1); 
}
int main() {
	scanf("%d%d%d",&n,&m,&r);  
	for(int i=1;i<=n;++i)	scanf("%d",w+i);
	int x,y;
	for(int i=1;i<=n-1;++i){
		scanf("%d%d",&x,&y);
		eg[x].pb(y),eg[y].pb(x);
	}
	dfs1(r,0); //两次从根dfs 
	dfs2(r,r);
	for(int i=1;i<=n;++i) add(id[i],id[i],nw[id[i]]);
	
	return 0;
}