文章標題 SPOJ-COT - Count on a tree（LCA+主席樹）

題目 鏈接

題意：求樹上的路徑 u->v上第k小的節點

分析：普通的第K大，當前的這顆線段樹是在前面一顆線段樹的基礎上建立的，而樹上的第K大，當前的線段樹可以在其父節點的線段樹建立起來。所以我們查詢u->v上的第k大就是rt[u]+rt[v]-rt[lca(u,v)]-rt[fa[lca(u,v)]]的第k大。

代碼：

#include <iostream>
#include <cstdio>
#include <cstring>
#include <string>
#include <queue>
#include <set>
#include <map>
#include <algorithm>
#include <math.h>
#include <vector>
using namespace std;
typedef long long ll;

const int mod=1e9+7;
const int maxn=2e5+10;
int n,q;
int a[maxn],ha[maxn]; 

int sum[maxn*40],ls[maxn*40],rs[maxn*40],rt[maxn*40];
int tot;
int sz;
//主席樹 
void build (int &num,int l,int r){
    num=++tot;
    sum[num]=0;
    if (r==l)return;
    int mid=(l+r)/2;
    build (ls[num],l,mid);
    build (rs[num],mid+1,r);
}

void update(int &num,int l,int r,int last,int pos){
    num=++tot;
    ls[num]=ls[last];
    rs[num]=rs[last];
    sum[num]=sum[last]+1;
    if (l==r)return ;
    int mid=(l+r)/2;
    if (pos<=mid)update(ls[num],l,mid,ls[last],pos);
    else update(rs[num],mid+1,r,rs[last],pos);
}

int query(int u,int v,int lca,int lcafa,int l,int r,int k){
    if (l==r)return l;
    int mid=(l+r)/2;
    int cnt=sum[ls[u]]+sum[ls[v]]-sum[ls[lca]]-sum[ls[lcafa]];
    if (cnt>=k)return query(ls[u],ls[v],ls[lca],ls[lcafa],l,mid,k);
    else return query(rs[u],rs[v],rs[lca],rs[lcafa],mid+1,r,k-cnt);
}

//LCA
int dis[maxn];//dis[i]表示節點i到樹根的距離 
int dp[maxn*2][20];
int first[maxn];//first[i]表示i這個節點第一次出現的標號 
int dep[maxn*2];//深度 
int point[maxn*2];//point[i]表示標號爲i所對應的節點 
int vis[maxn];//標記節點是否被訪問過了 
int cnt;//遍歷是節點的數目 
int fa[maxn];
vector<int>G[maxn];

void dfs(int u,int father,int idx){
    vis[u]=true;
    point[++cnt]=u;//point表示節點的標號
    first[u]=cnt;//first表示u這個節點第一次出現的位置
    dep[cnt]=idx;//深度
    fa[u]=father;
    update(rt[u],1,sz,rt[father],a[u]);
    for (int i=0;i<G[u].size();i++){
        int v=G[u][i];
        if (!vis[v]){
            dfs(v,u,idx+1);
            point[++cnt]=u;
            dep[cnt]=idx;
        }
    }
}
void RMQ_init(int n){
    for (int i=1;i<=n;i++)dp[i][0]=i;
    for (int k=1;(1<<k)<=n;k++){
        for (int i=1;i+(1<<k)-1<=n;i++){
            int a=dp[i][k-1];
            int b=dp[i+(1<<(k-1))][k-1];
            if (dep[a]>dep[b])dp[i][k]=b;
            else dp[i][k]=a;
        }
    }
} 

int RMQ(int l,int r){
    int k=0;
    while ((1<<(k+1))<=r-l+1)k++;
    int a=dp[l][k];
    int b=dp[r-(1<<k)+1][k];//保存的是編號
    return dep[a]>dep[b]?b:a;
}

int LCA(int u,int v){//求LCA 
    int x=first[u],y=first[v];
    if (x>y)swap(x,y);
    int ans=RMQ(x,y);
    return point[ans];
}

int main()
{
    while (scanf ("%d%d",&n,&q)!=EOF){
        for (int i=0;i<=n;i++)G[i].clear();
        memset (vis,0,sizeof (vis));
        for (int i=1;i<=n;i++){
            scanf ("%d",&a[i]);
            ha[i]=a[i];
        }
        int u,v,k;
        for (int i=0;i<n-1;i++){
            scanf ("%d%d",&u,&v);
            G[u].push_back(v);
            G[v].push_back(u);
        }
        sort(ha+1,ha+1+n);
        sz=unique(ha+1,ha+1+n)-(ha+1);
        for (int i=1;i<=n;i++){
            a[i]=lower_bound(ha+1,ha+1+sz,a[i])-ha;
        }
        tot=cnt=0;
        build(rt[0],1,sz);
        dfs(1,0,1);
        RMQ_init(2*n-1);
        while (q--){
            scanf ("%d%d%d",&u,&v,&k);
            int lca=LCA(u,v);
            int ans=query(rt[u],rt[v],rt[lca],rt[fa[lca]],1,sz,k);
            printf ("%d\n",ha[ans]);
        }
    }
    return 0;
}
/*
8 6
105 2 9 3 8 5 7 7
1 2        
1 3
1 4
3 5
3 6
3 7
4 8
2 5 1
2 5 2
2 5 3
2 5 4
7 8 2
5 7 1
*/