[HAOI2015]树上操作 树链剖分

[HAOI2015]树上操作 树链剖分

题面

裸树剖，线段树带lazy tag

#include <cstdio>
#include <algorithm>
#define MAXN 100010
#define ll long long
#define sl (x<<1)
#define sr (x<<1|1)
using namespace std;
int head[MAXN],vv[MAXN*2],nxt[MAXN*2],tot;
inline void add_edge(int u, int v){
	vv[++tot]=v;
	nxt[tot]=head[u];
	head[u]=tot;
}
int n;
ll w[MAXN];
int fa[MAXN],sz[MAXN],mxs[MAXN],dep[MAXN];
void dfs1(int u, int f){
	fa[u]=f;
	dep[u]=dep[f]+1;
	sz[u]=1;
	int mxsz=-1;
	for(int i=head[u];i;i=nxt[i]){
		int v=vv[i];
		if(v==f) continue;
		dfs1(v, u);
		sz[u]+=sz[v];
		if(mxsz<sz[v]){
			mxsz=sz[v];
			mxs[u]=v;
		}
	}
}
int topf[MAXN],idx[MAXN],cnt;
ll wnew[MAXN];
void dfs2(int u, int top){
	idx[u]=++cnt;
	wnew[cnt]=w[u];
	topf[u]=top;
	if(mxs[u]==0) return;
	dfs2(mxs[u], top);
	for(int i=head[u];i;i=nxt[i]){
		int v=vv[i];
		if(v==fa[u]||v==mxs[u]) continue;
		dfs2(v,v);
	}
}
struct nod{
	ll sum,lazy;
} tre[MAXN*4];
void push_up(int x){
	tre[x].sum=tre[sl].sum+tre[sr].sum;
}
void buildt(int x, int l, int r){
	if(l==r){
		tre[x].sum=wnew[l];
		return;
	}
	int mid=(l+r)>>1;
	buildt(sl, l, mid);
	buildt(sr, mid+1, r);
	push_up(x);
}
void push_down(int x, int l, int r){
	if(tre[x].lazy==0) return;
	tre[sl].lazy+=tre[x].lazy;
	tre[sr].lazy+=tre[x].lazy;
	int mid=(l+r)>>1;
	tre[sl].sum+=tre[x].lazy*(mid-l+1);
	tre[sr].sum+=tre[x].lazy*(r-(mid+1)+1);
	tre[x].lazy=0;
}
void change(int x, int l, int r, int cl, int cr, ll val){
	if(cl<=l&&r<=cr){
		tre[x].sum+=val*(r-l+1);
		tre[x].lazy+=val;
		return;
	}
	push_down(x, l, r);
	int mid=(l+r)>>1;
	if(cl<=mid) change(sl, l, mid, cl, cr, val);
	if(mid<cr) change(sr, mid+1, r, cl, cr, val);
	push_up(x);
}
ll query(int x, int l, int r, int ql, int qr){
	if(ql<=l&&r<=qr){
		return tre[x].sum;
	}
	push_down(x, l, r);
	int mid=(l+r)>>1;
	ll ans=0;
	if(ql<=mid) ans+=query(sl, l, mid, ql, qr);
	if(mid<qr) ans+=query(sr, mid+1, r, ql, qr);
	return ans;
}
ll tre_query(int a, int b){
	ll ans=0;
	while(topf[a]!=topf[b]){
		if(dep[topf[a]]<dep[topf[b]]) swap(a,b);
		ans+=query(1, 1, n, idx[topf[a]], idx[a]);
		a=fa[topf[a]];
	}
	if(dep[a]<dep[b]) swap(a,b);
	ans+=query(1, 1, n, idx[b], idx[a]);
	return ans;
}
void tre_change(int a, int val){
	change(1, 1, n, idx[a], idx[a]+sz[a]-1, val);
}
int main()
{
	scanf("%d", &n);
	int q;
	scanf("%d", &q);
	for(int i=1;i<=n;++i) scanf("%lld", &w[i]);
	for(int i=1;i<n;++i){
		int a,b;
		scanf("%d %d", &a, &b);
		add_edge(a,b);
		add_edge(b,a);
	}
	dfs1(1,1);
	dfs2(1,1);
	buildt(1, 1, n);
	while(q--){
		int type;
		scanf("%d", &type);
		int x;
		ll a;
		if(type==1){
			scanf("%d %lld", &x, &a);
			change(1, 1, n, idx[x], idx[x], a);
		}else if(type==2){
			scanf("%d %lld", &x, &a);
			tre_change(x, a);
		}else{
			scanf("%d", &x);
			printf("%lld
", tre_query(x, 1));
		}
	}
	return 0;
}