SP10628

SP10628 COT - Count on a tree

给定一棵有 \(n\) 个节点的数，每个点有权值，每次操作输出节点 \(u,v\) 之间的第 \(k\) 小值。

因为有静态区间第 \(k\) 小，所以考虑主席树。
因为题目要求在树上操作，所以考虑树上差分。
建树的过程就是跑一个 dfs，每个节点继承其父节点的权值建树。
对于一般的维护区间第 \(k\) 小的主席树，应用前缀和的思想，之后差分出 \(rt_r - rt_{l-1}\) 得到结果。
转成树上差分就变成 \(rt_u + rt_v - rt_{lca(u,v)} - rt_{fa_{lca(u,v)}}\)。
剩下的就是主席树的模板了。
具体实现看代码。

#include <bits/stdc++.h>
#define MID int m=(l+r)>>1;
#define V e[i].v
using namespace std;int rd(){
	int w=0,v=1;char c=getchar();while(c<'0'||c>'9'){if(c=='-')v=-1;c=getchar();}
	while(c>='0'&&c<='9'){w=(w<<1)+(w<<3)+(c&15);c=getchar();}return w*v;
}void wr(int x){
    char c[20];int l=0;if(x<0){putchar((1<<5)+(1<<3)+(1<<2)+1);x=~x+1;}
    do{c[l++]=x%10+(1<<4)+(1<<5); x/=10; }while(x>0);for(int i=l-1;i>=0;i--)putchar(c[i]);putchar('\n');
}const int N=2e5+100;int a[N],n,m,b[N],tot,lst,fa[N],tp,rt[N*18];struct tree{int l,r,z;}s[N*18];
struct E{int v,nt;}e[N<<1];int fir[N],c;void I(int u,int v){e[++c].v=v;e[c].nt=fir[u];fir[u]=c;}
struct ox{
	int rt=1,a,b,n,m,cnt,dep[N],dx[N],olx[N],st[N][25];void dfs(int u,int d){
		dx[u]=++cnt;olx[cnt]=u;dep[cnt]=d;
		for(int i=fir[u];i;i=e[i].nt)if(!dx[V])dfs(V,d+1),olx[++cnt]=u,dep[cnt]=d;
	}void init(){
		for(int i=1;i<=cnt;i++)st[i][0]=i;for(int j=1;j<=log2(cnt);j++)for(int i=1;i+(1<<j)-1<=cnt;i++)
		a=st[i][j-1],b=st[i+(1<<(j-1))][j-1],st[i][j]=(dep[a]<dep[b])?a:b;
	}int lca(int l,int r){
		l=dx[l],r=dx[r];if(l>r)swap(l,r);int k=log2(r-l+1);a=st[l][k],b=st[r-(1<<k)+1][k];
		return (dep[a]<dep[b])?olx[a]:olx[b];
	}
}LCA;int id(int x){return lower_bound(a+1,a+tot+1,x)-a;}int C(int t,int l,int r,int x){
	int p=++tp;s[p]=s[t];s[p].z++;if(l==r)return p;MID 
	if(x<=m)s[p].l=C(s[t].l,l,m,x);else s[p].r=C(s[t].r,m+1,r,x);return p;
}void d2(int u,int f){fa[u]=f;rt[u]=C(rt[fa[u]],1,N,id(b[u]));for(int i=fir[u];i;i=e[i].nt)if(V!=f)d2(V,u);}
int query(int a,int b,int c,int d,int l,int r,int k){
	if(l==r)return l;int e=s[s[a].l].z+s[s[b].l].z-s[s[c].l].z-s[s[d].l].z;
	MID if(e>=k)return query(s[a].l,s[b].l,s[c].l,s[d].l,l,m,k);
	else return query(s[a].r,s[b].r,s[c].r,s[d].r,m+1,r,k-e);
}int main(){
	n=rd();m=rd();for(int i=1;i<=n;i++)b[i]=rd(),a[i]=b[i];sort(a+1,a+n+1);tot=unique(a+1,a+n+1)-a-1;
	for(int i=1,u,v;i<n;i++)u=rd(),v=rd(),I(u,v),I(v,u);LCA.dfs(1,1);LCA.init();d2(1,0);
	for(int i=1,u,v,k,l;i<=m;i++)u=rd(),v=rd(),k=rd(),l=LCA.lca(u,v),wr(a[query(rt[u],rt[v],rt[l],rt[fa[l]],1,N,k)]);
}