SPOJ COT Count on a tree（树上路径第k小 主席树）

题目链接

题意：

求树上A，B两点路径上第K小的数

分析：

同样是可持久化线段树，只是这一次我们用它来维护树上的信息。

我们之前已经知道，可持久化线段树实际上是维护的一个前缀和，而前缀和不一定要出现在一个线性表上。

比如说我们从一棵树的根节点进行DFS，得到根节点到各节点的距离dist[x]——这是一个根-x路径上点与根节点距离的前缀和。

利用这个前缀和，我们可以解决一些树上任意路径的问题，比如在线询问[a,b]点对的距离——答案自然是dist[a]+dist[b]-2*dist[lca(a,b)]。

同理，我们可以利用可持久化线段树来解决树上任意路径的问题。

DFS遍历整棵树，然后在每个节点上建立一棵线段树，某一棵线段树的“前一版本”是位于该节点父亲节点fa的线段树。

利用与之前类似的方法插入点权(排序离散)。那么对于询问[a,b]，答案就是root[a]+root[b]-root[lca(a,b)]-root[fa[lca(a,b)]]上的第k大。

//#pragma comment(linker, "/STACK:1024000000,1024000000")
#include<iostream>
#include<stdio.h>
#include<math.h>
#include <string>
#include<string.h>
#include<map>
#include<queue>
#include<set>
#include<utility>
#include<vector>
#include<algorithm>
#include<stdlib.h>
using namespace std;
#define eps 1e-8
#define pii pair<int,int>
#define inf 0x3f3f3f3f
#define rd(x) scanf("%d",&x)
#define rd2(x,y) scanf("%d%d",&x,&y)
#define ll long long int
#define mod 1000000007
#define maxn 100050
#define maxm 5000005
int ls[maxm],rs[maxm],sum[maxm];
int tot,pos,n,m;
int a[maxn],f[maxn],rt[maxn],faa[maxn];
int nn,u,v,root,k;
struct node{
    int v,next;
}edge[maxn*2];
int head[maxn];
void addedge(int u,int v){
    edge[++tot].v=v;edge[tot].next=head[u];head[u]=tot;
}
void update(int &x,int pr,int l,int r,int p){//对于求第i个线段树可以理解为，
    x=++pos;sum[x]=sum[pr]+1;        //先复制第i-1个，相同的位置公用空间，不同重新开辟
    if(l==r) return;
    ls[x]=ls[pr];rs[x]=rs[pr];
    int mid=(l+r)>>1;
    if(p<=mid) update(ls[x],ls[pr],l,mid,p);
    else update(rs[x],rs[pr],mid+1,r,p);
}
void build(int &x,int l,int r){
    x=++pos;sum[x]=0;
    if(l==r) return;
    int mid=(l+r)/2;
    build(ls[x],l,mid);
    build(rs[x],mid+1,r);
}
int h(int x){//获取值x在线段树的位置
    return lower_bound(f+1,f+1+nn,x)-f;
}
int query(int L,int R,int k,int lc,int flc,int l,int r){
    if(l==r) return f[l];
    int tmp=sum[ls[R]]+sum[ls[L]]-sum[ls[lc]]-sum[ls[flc]];
    int mid=(l+r)>>1;
    if(k<=tmp) return query(ls[L],ls[R],k,ls[lc],ls[flc],l,mid);
    else return query(rs[L],rs[R],k-tmp,rs[lc],rs[flc],mid+1,r);
}
int rmq[2*maxn];//求公共祖先在线算法
struct ST{
    int mm[2*maxn];
    int dp[2*maxn][20];
    void init(int n){
        mm[0]=-1;
        for(int i=1;i<=n;i++){
            mm[i]=((i&(i-1))==0)?mm[i-1]+1:mm[i-1];
            dp[i][0]=i;
        }
        for(int j=1;j<=mm[n];j++)
            for(int i=1;i+(1<<j)-1<=n;i++)
            dp[i][j]=rmq[dp[i][j-1]]<rmq[dp[i+(1<<(j-1))][j-1]]?
            dp[i][j-1]:dp[i+(1<<(j-1))][j-1];
    }
    int query(int a,int b){
        if(a>b) swap(a,b);
        int k=mm[b-a+1];
        return rmq[dp[a][k]]<=rmq[dp[b-(1<<k)+1][k]]?dp[a][k]:dp[b-(1<<k)+1][k];
    }
}st;
int F[maxn*2],P[maxn],cnt;
void dfs(int x,int fa,int deep){
    faa[x]=fa;
    F[++cnt]=x;rmq[cnt]=deep;P[x]=cnt;
    update(rt[x],rt[fa],1,nn,h(a[x]));
    for(int i=head[x];i!=-1;i=edge[i].next){
        if(edge[i].v==fa) continue;
        dfs(edge[i].v,x,deep+1);
        F[++cnt]=x;rmq[cnt]=deep;
    }
}
int query_lca(int u,int v){
    return F[st.query(P[u],P[v])];
}
int main()
{
    rd2(n,m);
    for(int i=1;i<=n;i++) {
            rd(a[i]);f[i]=a[i];
    }
    sort(f+1,f+1+n);
    nn=unique(f+1,f+1+n)-f-1;//离散化线段树
    memset(head,-1,sizeof(head));
    for(int i=1;i<n;i++){
        rd2(u,v);
        addedge(u,v);
        addedge(v,u);
    }
    pos=tot=cnt=0;
    build(rt[0],1,nn);
    dfs(1,0,0);//构建线段树组，根据父子关系
    st.init(2*n-1);
    for(int i=1;i<=m;i++){
        scanf("%d%d%d",&u,&v,&k);
        int lc=query_lca(u,v);
        printf("%d\n",query(rt[u],rt[v],k,rt[lc],rt[faa[lc]],1,nn));
    }
    return 0;
}