樹鏈剖分模板

樹鏈剖分:樹上統計的一種優秀算法,基本思想是用將樹切割成幾段長鏈再配合數據結構轉爲區間問題。

我使用了線段樹,需要注意的是樹鏈剖分的額外空間極大,需要儘可能壓縮

線段樹的部分並無改動。

關鍵在於剖分的過程。

分兩次dfs完成

第一遍求出每顆子樹大小;每個節點的父親節點、深度、重兒子(即子樹大小最大的兒子)

第二遍按照重兒子優先的原則將節點重新編號,分成長鏈,這樣的做法可以使每條鏈上的序號連續,而深搜使各子樹中的節點集合也爲連續區間。建立線段樹

 操作:

如果操作某節點x子樹上的關鍵值,較簡單,操作序號id[x]~id[x]+size[x]-1(易證)

如果操作兩點間路徑上的關鍵值,稍複雜,每次將兩點所在鏈分別處理直到兩點位於同一鏈即可

下面是程序

#include<iostream>
#include<cstdio>
#include<cstring>
#define MAXN 100001
using namespace std;
int n,m,r,p;
int a[MAXN];
int h,head[MAXN];
int fa[MAXN],size[MAXN],dep[MAXN],hson[MAXN];
int id[MAXN],top[MAXN],tot;
struct Edge
{
	int v,next;
}edge[MAXN*2];
struct node
{
	int lazy,sum;
}tr[MAXN*4];
void change(int o,int nl,int nr,int l,int r,int a)
{
	if(nl==l&&nr==r)
	{
	    tr[o].lazy+=a;
		tr[o].sum+=a*(nr-nl+1);
		return;
	}
	int mid=(nl+nr)>>1;
	change(o<<1,nl,mid,nl,mid,tr[o].lazy);change(o<<1|1,mid+1,nr,mid+1,nr,tr[o].lazy);
	tr[o].lazy=0;
	if(r<=mid) change(o<<1,nl,mid,l,r,a);
	else if(l>mid) change(o<<1|1,mid+1,nr,l,r,a);
	else change(o<<1,nl,mid,l,mid,a),change(o<<1|1,mid+1,nr,mid+1,r,a);
	tr[o].sum=(tr[o<<1].sum+tr[o<<1|1].sum)%p;
}
int ask(int o,int nl,int nr,int l,int r)
{
	if(l==nl&&r==nr)
	{
		return tr[o].sum;
	}
	int mid=(nl+nr)>>1;
	change(o<<1,nl,mid,nl,mid,tr[o].lazy);change(o<<1|1,mid+1,nr,mid+1,nr,tr[o].lazy);
	tr[o].lazy=0;
	if(r<=mid) return ask(o<<1,nl,mid,l,r);
	else if(l>mid) return ask(o<<1|1,mid+1,nr,l,r);
	else return ask(o<<1,nl,mid,l,mid)+ask(o<<1|1,mid+1,nr,mid+1,r);
}
void add(int u,int v)
{
	h++;
	edge[h].v=v;
	edge[h].next=head[u];
	head[u]=h;
}
void dfs1(int last,int u)
{
	fa[u]=last;
	dep[u]=dep[last]+1;
	size[u]=1;
	int maxx=-1;
	for(int i=head[u];i;i=edge[i].next)
	{
		int v=edge[i].v;
		if(v==last) continue;
		dfs1(u,v);
		size[u]+=size[v];
		if(size[v]>maxx) maxx=size[v],hson[u]=v;
	}
}
void dfs2(int u,int topf)
{
	id[u]=++tot;
    top[u]=topf;
    change(1,1,n,tot,tot,a[u]);
    if(!hson[u]) return;
    dfs2(hson[u],topf);
    for(int i=head[u];i;i=edge[i].next)
    {
    	int v=edge[i].v;
    	if(v==fa[u]||v==hson[u]) continue;
    	dfs2(v,v);
	}
}
int ask_tree1(int x,int y)
{
	int res=0;
	while(top[x]!=top[y])
	{
		if(dep[top[x]]<dep[top[y]]) swap(x,y);
		res=(res+ask(1,1,n,id[top[x]],id[x]))%p;
		x=fa[top[x]];
	}
	if(dep[x]>dep[y]) swap(x,y);
	res=(res+ask(1,1,n,id[x],id[y]))%p;
	return res;
}
int ask_tree2(int x)
{
	return ask(1,1,n,id[x],id[x]+size[x]-1);
}
void change_tree1(int x,int y,int a)
{
	while(top[x]!=top[y])
	{
		if(dep[top[x]]<dep[top[y]]) swap(x,y);
		change(1,1,n,id[top[x]],id[x],a);
		x=fa[top[x]];
	}
	if(dep[x]>dep[y]) swap(x,y);
	change(1,1,n,id[x],id[y],a);
} 
void change_tree2(int x,int a)
{
	change(1,1,n,id[x],id[x]+size[x]-1,a);
}
int main()
{
	cin>>n>>m>>r>>p;
	for(int i=1;i<=n;i++) scanf("%d",&a[i]);
	for(int i=1;i<n;i++)
	{
		int a,b;
		scanf("%d%d",&a,&b);
		add(a,b);add(b,a);
	}
	dfs1(0,r);dfs2(r,r);
	for(int i=1;i<=m;i++)
	{
		int op,x,y,z;
		scanf("%d",&op);
		if(op==1)
		{
			scanf("%d%d%d",&x,&y,&z);
			change_tree1(x,y,z);
		}
		if(op==2)
		{
			scanf("%d%d",&x,&y);
			printf("%d\n",ask_tree1(x,y));
		}
		if(op==3)
		{
			scanf("%d%d",&x,&y);
			change_tree2(x,y);
		}
		if(op==4)
		{
			scanf("%d",&x);
			printf("%d\n",ask_tree2(x));
		}
	}
}

 

發佈了6 篇原創文章 · 獲贊 5 · 訪問量 537
發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章