洛谷 P3178 [HAOI2015]树上操作

版权声明：欢迎转载欢迎评论！ https://blog.csdn.net/rabbit_ZAR/article/details/85118382

题目：树上操作

思路：树剖模板题。

代码：

#include<bits/stdc++.h>
using namespace std;

#define read(x) scanf("%lld",&x)
#define maxn 100000
#define ll long long

struct Node {
	ll lstid;
	ll w,fa;
	ll sz,d,hson;
	ll tp;
};

ll n,m;
ll rt;

Node a[maxn+5];
vector<ll> tr[maxn+5];

ll cnt=0;
ll frtid[maxn+5];

struct segTree {
	#define lson (o*2)
	#define rson (o*2+1)
	#define mid (L+(R-L)/2)
	
	ll c[maxn*10+5],lzy[maxn*10+5];
	ll len[maxn*10+5];
	
	void push_up(ll o) {
		c[o]=c[lson]+c[rson];
	}
	
	void build(ll o,ll L,ll R) {
		len[o]=R-L+1;
		if(L==R) {
			c[o]=a[frtid[L]].w;
			return ;
		}
		build(lson,L,mid),build(rson,mid+1,R);
		push_up(o);
	}
	
	ll P,Q;
	
	void push_down(ll o) {
		lzy[lson]=lzy[lson]+lzy[o],lzy[rson]=lzy[rson]+lzy[o];
		c[lson]=c[lson]+lzy[o]*len[lson],c[rson]=c[rson]+lzy[o]*len[rson];
		lzy[o]=0;
	}
	
	ll Qry(ll o,ll L,ll R) {
		if(L>Q||R<P) return 0;
		if(L>=P&&R<=Q) return c[o];
		push_down(o);
		ll ans=Qry(lson,L,mid)+Qry(rson,mid+1,R);
		return ans;
	}
	
	ll Query(ll x,ll y) {
		if(x>y) swap(x,y);
		P=x,Q=y;
		return Qry(1,1,n);
	}
	
	void Upd(ll o,ll L,ll R,ll d) {
		if(L>Q||R<P) return ;
		if(L>=P&&R<=Q) {
			c[o]=c[o]+d*len[o];
			lzy[o]=d+lzy[o];
			return ;
		}
		push_down(o);
		Upd(lson,L,mid,d),Upd(rson,mid+1,R,d);
		push_up(o);
	}
	
	void Update(ll x,ll y,ll d) {
		if(x>y) swap(x,y);
		P=x,Q=y;
		Upd(1,1,n,d);
	}
};

segTree sgtr;

struct HLD{
	void update(ll x,ll y,ll z) {
		while(a[x].tp!=a[y].tp) {
			if(a[a[x].tp].d<a[a[y].tp].d) swap(x,y);
			sgtr.Update(a[x].lstid,a[a[x].tp].lstid,z);
			x=a[a[x].tp].fa;
		}
		sgtr.Update(a[x].lstid,a[y].lstid,z);
	}
	
	ll query(ll x,ll y) {
		ll ans=0;
		while(a[x].tp!=a[y].tp) {
			if(a[a[x].tp].d<a[a[y].tp].d) swap(x,y);
			ans=ans+ sgtr.Query(a[x].lstid,a[a[x].tp].lstid);
			x=a[a[x].tp].fa;
		}
		ans=ans+ sgtr.Query(a[x].lstid,a[y].lstid);
		return ans;
	}
};

HLD hld;

void readin() {
	read(n),read(m);
	for(ll i=1;i<=n;i++) read(a[i].w);
	for(ll i=1;i<n;i++) {
		ll x,y;
		read(x),read(y);
		tr[x].push_back(y);
		tr[y].push_back(x);
	}
}

void dfs1(ll x,ll fa) {
	a[x].fa=fa;
	a[x].sz=1,a[x].d=a[fa].d+1;
	ll maxson=0,maxsz=-1;
	for(ll i=0;i<tr[x].size();i++) {
		ll y=tr[x][i];
		if(y==fa) continue;
		dfs1(y,x);
		a[x].sz+=a[y].sz;
		if(a[y].sz>maxsz) {
			maxsz=a[y].sz;
			maxson=y;
		}
	}
	a[x].hson=maxson;
}

void dfs2(ll x,ll tp,ll fa) {
	a[x].tp=tp;
	a[x].lstid=++cnt,frtid[cnt]=x;
	if(a[x].hson==0) return;
	dfs2(a[x].hson,tp,x);
	for(ll i=0;i<tr[x].size();i++) {
		ll y=tr[x][i];
		if(y==fa||y==a[x].hson) continue;
		dfs2(y,y,x);
	}
}

int main() {
	readin();
	rt=1;
	dfs1(rt,0);
	dfs2(rt,rt,0);
	
	sgtr.build(1,1,n);
	while(m--) {
		ll opr;
		read(opr);
		if(opr==1) {
			ll x,z;
			read(x),read(z);
			sgtr.Update(a[x].lstid,a[x].lstid,z);
		} else if(opr==2) {
			ll x,z;
			read(x),read(z);
			sgtr.Update(a[x].lstid,a[x].lstid+a[x].sz-1,z);
		} else  {
			ll x;
			read(x);
			printf("%lld\n",hld.query(rt,x));
		}
	}
	
	return 0;
}