P4175 [CTSC2008]网络管理

题目大意

给出一棵树,每个点有一个权值,支持查询树链上第 $k$ 大数和单点修改.

分析

先把修改操作去掉开始(题目),不带修改树链第 $k$ 大,要查询第 $k$ 大就很容易想到用一些数据结构维护(废话),例如权值线段树,平衡树等等.

先从一个看似不相关的问题开始考虑,计算一条树链上所有数的和需要怎么办?

很显然这个东西可以树剖一下,再用前缀和优化,可以做到 $\mathcal{O}(lo{g}_{2}N)$ 的时间复杂度单词查询.

用这个方法做这道题可以想到树剖之后用主席树维护在DFS序中前缀每个数出现次数,再可以计算区间每个数出现次数,找出这条树链中的的每一段计算的区间的开和结尾的主席树根节点编号,再通过类似前缀和计算区间和的方法计算,在树上二分查询第 $k$ 大数,查询时间复杂度变为 $\mathcal{O}(lo{g}_2^{2}N)$.

但是上面这个方法并不是很优秀,可以再结合容斥的思想,$su{m}_i$ 表示第 $i$ 个节点到根节点这一条链上所有数的和,那么一条 $u$ 到 $v$ 树链的和就可以表示为 $su{m}_u+su{m}_v-su{m}_{LCA(u,v)}\ast 2+va{l}_{LCA(u,v)}$,其中 $-su{m}_{LCA(u,v)}+va{l}_{LCA(u,v)}$ 可以表示为 $su{m}_{fathe{r}_{LCA(u,v)}}$.

通过上面这个方法可以想到用主席树维护每个节点到根节点这一条链上面每个数出现的次数,最后只需要将这样四颗主席树进行加减就可以直接在树上二分,时间复杂度是 $\mathcal{O}(lo{g}_{2}N)$ 的.

再回到本题,本题有修改,但是主席树显然是没法直接做到修改的.那么可以从每个修改所造成的影响来看,仍然用权值线段树维护每个节点到根节点这一条树链上每个数出现的次数,那么如果一个点的值消失会造成的影响只会在以这个点为根节点的子树中,它的每一个子节点(包括自己)的权值线段树中减去 $va{l}_{now}$,添加一个数同理加上.

那么问题就变成了如何给一颗子树加上一个值,在DFS序中,每一颗子树的节点的编号都是连续的,所以就变成了区间加,然后树状数组维护一下就好了.

代码

#include<bits/stdc++.h>
#define REP(i,first,last) for(int i=first;i<=last;++i)
#define DOW(i,first,last) for(int i=first;i>=last;--i)
using namespace std;
const int MAXN=4e5+7;
const int MAX_NUM=1e8;
int N,M,val[MAXN];
int edge_head[MAXN];
int edge_cnt=0;
struct Edge
{
	int to,next;
}edge[MAXN];
#define FOR(now) for(int i_edge=edge_head[now];i_edge;i_edge=edge[i_edge].next)
#define TO edge[i_edge].to
void AddEdge(int form,int to)
{
	edge[++edge_cnt].to=to;
	edge[edge_cnt].next=edge_head[form];
	edge_head[form]=edge_cnt;
}

int kth_father[MAXN][25];//比较懒,直接用了倍增LCA,一个裸的LCA,如果不
bool visit[MAXN];
int deep[MAXN],father[MAXN];
int tree_size[MAXN];
int dfs[MAXN],dfs_cnt=0;
int id[MAXN];
void DFS(int now)
{
	deep[now]=deep[father[now]]+1;
	kth_father[now][0]=father[now];
	tree_size[now]=1;//需要记录子树大小
	dfs[++dfs_cnt]=now;//记录下dfs序
	id[now]=dfs_cnt;
	REP(i,0,22)
	{
		kth_father[now][i+1]=kth_father[kth_father[now][i]][i];
	}
	FOR(now)
	{
		if(father[now]!=TO)
		{
			father[TO]=now;
			DFS(TO);
			tree_size[now]+=tree_size[TO];
		}
	}
}
int LCA(int x,int y)
{
	if(deep[x]<deep[y])
	{
		swap(x,y);
	}
	DOW(i,22,0)
	{
		if(deep[kth_father[x][i]]>=deep[y])
		{
			x=kth_father[x][i];
		}
		if(x==y)
		{
			return x;
		}
	}
	DOW(i,22,0)
	{
		if(kth_father[x][i]!=kth_father[y][i])
		{
			x=kth_father[x][i];
			y=kth_father[y][i];
		}
	}
	return father[x];
}
struct SegmentTree//线段树部分
{
	int sum,lson,rson;
}sgt[MAXN*32];
int sgt_cnt=0;
int root[MAXN];
int Lowbit(int now)
{
	return now&-now;
}
#define LSON sgt[now].lson
#define RSON sgt[now].rson
#define MIDDLE ((left+right)>>1)
#define LEFT LSON,left,MIDDLE
#define RIGHT RSON,MIDDLE+1,right
void PushUp(int now)
{
	sgt[now].sum=sgt[LSON].sum+sgt[RSON].sum;
}
void Updata(int num,int val,int &now,int left=1,int right=MAX_NUM)//修改操作不多说了
{
	if(num<left||right<num)
	{
		return;
	}
	if(!now)
	{
		now=++sgt_cnt;
	}
	if(left==right)
	{
		sgt[now].sum+=val;
		return;
	}
	Updata(num,val,LEFT);
	Updata(num,val,RIGHT);
	PushUp(now);
}
void Change(int p,int w,int check=1)//替换操作
{
	int left=id[p],right=id[p]+tree_size[p]-1;//计算开始和结束位置
	if(check)//开始加数时不需要删数
	{
		for(int i=left;i<=N;i+=Lowbit(i))//将这个数删去,同区间修改树状数组
		{
			Updata(val[p],-1,root[i]);
		}
		for(int i=right+1;i<=N;i+=Lowbit(i))
		{
			Updata(val[p],1,root[i]);
		}
		val[p]=w;
	}
	for(int i=left;i<=N;i+=Lowbit(i))//同理加上新的数
	{
		Updata(val[p],1,root[i]);
	}
	for(int i=right+1;i<=N;i+=Lowbit(i))
	{
		Updata(val[p],-1,root[i]);
	}
}
int add_cnt;
int dec_cnt;
int add_root[MAXN];//记录加上/删掉的树的当前节点
int dec_root[MAXN];
int GetSum()//计算当前范围的数的个数
{
	int result=0;
	REP(i,1,add_cnt)
	{
		result+=sgt[add_root[i]].sum;
	}
	REP(i,1,dec_cnt)
	{
		result-=sgt[dec_root[i]].sum;
	}
	return result;
}
int GetSumRight()//计算右子树的数的个数
{
	int result=0;
	REP(i,1,add_cnt)
	{
		result+=sgt[sgt[add_root[i]].rson].sum;
	}
	REP(i,1,dec_cnt)
	{
		result-=sgt[sgt[dec_root[i]].rson].sum;
	}
	return result;
}
void GetRootLeft()//将当前节点变为右子节点
{
	REP(i,1,add_cnt)
	{
		add_root[i]=sgt[add_root[i]].lson;
	}
	REP(i,1,dec_cnt)
	{
		dec_root[i]=sgt[dec_root[i]].lson;
	}
}
void GetRootRight()//变为左子节点
{
	REP(i,1,add_cnt)
	{
		add_root[i]=sgt[add_root[i]].rson;
	}
	REP(i,1,dec_cnt)
	{
		dec_root[i]=sgt[dec_root[i]].rson;
	}
}
int QueryKth(int k,int left=1,int right=MAX_NUM)//查询部分
{
	if(left==right)
	{
		return left;
	}
	int sum=GetSumRight();
	if(sum>=k)//如果右边够就查询右边
	{
		GetRootRight();
		return QueryKth(k,MIDDLE+1,right);
	}
	GetRootLeft();
	return QueryKth(k-sum/*需要减去右子树中数的个数*/,left,MIDDLE);
}
void Kth(int u,int v,int k)
{
	int lca=LCA(u,v);
	//按公式计算
	int add_1=id[u];
	int add_2=id[v];
	int dec_1=id[lca];
	int dec_2=id[father[lca]];
	add_cnt=0;
	dec_cnt=0;
	for(int i=add_1;i;i-=Lowbit(i))
	{
		add_root[++add_cnt]=root[i];
	}
	for(int i=add_2;i;i-=Lowbit(i))
	{
		add_root[++add_cnt]=root[i];
	}
	for(int i=dec_1;i;i-=Lowbit(i))
	{
		dec_root[++dec_cnt]=root[i];
	}
	for(int i=dec_2;i;i-=Lowbit(i))
	{
		dec_root[++dec_cnt]=root[i];
	}
	if(GetSum()<k)//如果总共也没有k就是不存在
	{
		printf("invalid request!\n");
		return;
	}
	printf("%d\n",QueryKth(k));
}
int main()
{
	scanf("%d%d",&N,&M);
	REP(i,1,N)
	{
		scanf("%d",&val[i]);
	}
	int fa,son;
	REP(i,1,N-1)
	{
		scanf("%d%d",&fa,&son);
		AddEdge(fa,son);
		AddEdge(son,fa);
	}
	DFS(1);//处理树上的信息
	REP(i,1,N)//把数加入
	{
		Change(i,val[i],0);
	}
	int k,x,y;
	REP(i,1,M)
	{
		scanf("%d%d%d",&k,&x,&y);
		if(k==0)
		{
			Change(x,y);//修改
		}
		if(k)
		{
			Kth(x,y,k);//查询
		}
	}
	return 0;
}