「SP10628 COT - Count on a tree」

主席树的综合运用题.

前置芝士

1. 可持久化线段树:其实就是主席树了.
2. LCA:最近公共祖先,本题需要在\(\log_2N\)及以内的时间复杂度内解决这个问题.

具体做法

主席树维护每个点到根节点这一条链上不同树出现的次数,然后发现这个东西是可以相减的,于是这条链上每个数出现的次数就变成了\(sum[u]+sum[v]-2*sum[LCA(u,v)]\).然后就可以发现这个是错的,如果按这个式子计算最后的链上就没有LCA位置的值了,所以在范围包含\(val[LCA(u,v)]\)时需要加一,然后正常去找kth就好了.

代码

#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=2e5+7;

int N,M,point[maxN];
int root[maxN];
int sor[maxN];
map<int,int>Hash;
int val[maxN];
//一下为主席树部分
int point_cnt=0;
struct Tree//对于树的一个结构体
{
	int sum,lson,rson;
}tree[maxN*32];
//主席树标准define
#define LSON tree[now].lson
#define RSON tree[now].rson
#define MIDDLE ((left+right)>>1)
#define LEFT LSON,left,MIDDLE
#define RIGHT RSON,MIDDLE+1,right
#define NEW_LSON tree[new_tree].lson
#define NEW_RSON tree[new_tree].rson
void PushUp(int now)//合并信息
{
	tree[now].sum=tree[LSON].sum+tree[RSON].sum;
}
void UpData(int num,int &new_tree,int now,int left=1,int right=N)
//修改并产生一个新版本到new_tree中
{
	if(num<left||right<num)
	{
		new_tree=now;
		return;
	}
	new_tree=++point_cnt;
	if(left==right)
	{
		tree[new_tree].sum=tree[now].sum+1;
		return;
	}
	UpData(num,NEW_LSON,LEFT);
	UpData(num,NEW_RSON,RIGHT);
	PushUp(new_tree);
}
int QueryKth(int k,int LCA_num,int tree1/*相对于式子中的u*/,int tree2/*相当于式子中的v*/,int tree3/*相当于式子中的LCA(u,v)*/,int left=1,int right=N)
{
	if(left==right)
	{
		return left;
	}
	int sum=tree[tree[tree1].lson].sum
		   +tree[tree[tree2].lson].sum
		 -2*tree[tree[tree3].lson].sum;//计算出这条链上在left~right中的数出现的次数
	if(left<=LCA_num&&LCA_num<=MIDDLE)//如果LCA位置的数在当前的范围内就加一
	{
		sum++;
	}
	if(sum>=k)//左子树查询kth
	{
		return
		QueryKth(k,LCA_num,
		tree[tree1].lson,
		tree[tree2].lson,
		tree[tree3].lson,
		left,MIDDLE);
	}
	return//右子树查询
	QueryKth(k-sum,LCA_num,
	tree[tree1].rson,
	tree[tree2].rson,
	tree[tree3].rson,
	MIDDLE+1,right);
}
//一下为链式前向星部分(非主要内容,不多讲)
int head[maxN];
int edge_cnt=0;
struct Edge
{
	int to,next;
}edge[maxN];
#define FOR(now) for(int i_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=head[form];
	head[form]=edge_cnt;
}
//一下为倍增LCA部分(非主要内容,不多讲)
int kth_father[maxN][25];
bool visit[maxN];
int deep[maxN],father[maxN];
void DFS(int now)
{
	UpData(Hash[point[now]],root[now],root[father[now]]);//在DFS的过程中把主席树建好
	deep[now]=deep[father[now]]+1;
	kth_father[now][0]=father[now];
	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);
		}
	}
}
int LCA(int x,int y)//LCA核心代码
{
	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];
}

int main()
{
	scanf("%d%d",&N,&M);
	REP(i,1,N)
	{
		scanf("%d",&point[i]);
		sor[i]=point[i];
	}
	int fa,son;
	REP(i,1,N-1)
	{
		scanf("%d%d",&fa,&son);
		AddEdge(fa,son);
		AddEdge(son,fa);
	}
	sort(sor+1,sor+1+N);//离散化
	sor[0]=-114514;
	int cnt_num=0;
	REP(i,1,N)
	{
		if(sor[i]!=sor[i-1])
		{
			Hash[sor[i]]=++cnt_num;
			val[cnt_num]=sor[i];
		}
	}
	N=cnt_num;
	DFS(1);//DFS,建树+倍增LCA预处理
	int u,v,k;
	int point_LCA;
	REP(i,1,M)
	{
		scanf("%d%d%d",&u,&v,&k);
		point_LCA=LCA(u,v);//找到LCA
		printf("%d\n",
			val[QueryKth(k,Hash[point[point_LCA]],
					 root[u],root[v],root[point_LCA])]
			);//直接查找kth
	}
	return 0;
}