SPOJ Count on a tree 在線LCA + 主席樹

題意：求樹上a ， b路徑上第k大的點
主席樹a + 主席樹b - 主席樹（LCA（a，b）） - 主席樹（FA（LCA（a，b）））

代碼：

#include <bits/stdc++.h>
#define sf scanf
#define pf printf
using namespace std;
const int maxn = 100000 + 50;
/** 主席樹 */
const int tot_seg = maxn * 60;
int ch[tot_seg][2],SUM[tot_seg];
int tot;
#define lson(rt) (ch[rt][0])
#define rson(rt) (ch[rt][1])
void Insert(int prt,int rt,int l,int r,int x){
    SUM[rt] = SUM[prt] + 1;
    while(l < r){
        int mid = l + r >> 1;
        int kind = (x > mid);
        ch[rt][kind] = tot++;
        ch[rt][!kind] = ch[prt][!kind];
        rt = ch[rt][kind],prt = ch[prt][kind];
        SUM[rt] = SUM[prt] + 1;
        if(!kind) r = mid;
        else l = mid + 1;
    }
}
int Search(int art,int brt,int crt,int drt,int l,int r,int k){
    while(l < r){
        int mid = l + r >> 1,cnt = SUM[lson(art)] + SUM[lson(brt)] - SUM[lson(crt)] - SUM[lson(drt)];
        if(k <= cnt){
            r = mid;
            art = lson(art),brt = lson(brt),crt = lson(crt),drt = lson(drt);
        }
        else{
            l = mid + 1,k = k - cnt;
            art = rson(art),brt = rson(brt),crt = rson(crt),drt = rson(drt);
        }
    }
    return l;
}
void Tree_Init(){
    tot = 1;ch[0][0] = ch[0][1] = SUM[0] = 0;
}



int rts[maxn],value[maxn],Ar[maxn];
map<int ,int > HASH;

/** LCA */
vector<int> Adj[maxn];
int dep[maxn],R[maxn],dp[maxn * 2][32],vis_time,FA[maxn];
void DFS(int u,int fa,int Ar_len){
    Insert(rts[fa],rts[u] = tot++,1,Ar_len,value[u]);
    dep[u] = dep[FA[u] = fa] + 1;
    dp[R[u] = vis_time][0] = u;
    vis_time++;
    int len = Adj[u].size();
    for(int i = 0;i < len;++i){
        int v = Adj[u][i];
        if(v == fa) continue;
        DFS(v,u,Ar_len);dp[vis_time++][0] = u;
    }
}
void Init_ST(){
    for(int len = 1;1 << len <= vis_time;++len){
        for(int i = 0;(i + (1 << len)) <= vis_time;++i){
            int a = dp[i][len - 1],b = dp[i + (1 << len - 1)][len - 1];
            dp[i][len] = dep[a] < dep[b] ? a : b;
        }
    }
}
int RMQ(int a,int b){
    a = R[a],b = R[b];
    if(a > b) swap(a,b);
    int k = 0;
    while((1 << k + 1) <= (b - a + 1)) k++;
    a = dp[a][k],b = dp[b - (1 << k) + 1][k];
    return dep[a] < dep[b] ? a : b;
}


void ALL_INIT(int n){
    Tree_Init();
    HASH.clear();
    rts[0] = 0;
    vis_time = 0;
    for(int i = 0;i <= n;++i){
        Adj[i].clear();
    }
}
int main(){
//    freopen("read.txt","r",stdin);
    int n,m;
    while(~sf("%d %d",&n,&m)){
        ALL_INIT(n);
        for(int i = 1;i <= n;++i){
            sf("%d",&value[i]);
            Ar[i] = value[i];
        }
        sort(Ar + 1,Ar + 1 + n);
        int len = unique(Ar + 1,Ar + 1 + n) - Ar;
        for(int i = 1;i < len;++i) HASH[Ar[i]] = i;
        for(int i = 1;i <= n;++i) value[i] = HASH[value[i]];
        for(int i = 0;i < n - 1;++i){
            int a,b;sf("%d %d",&a,&b);
            Adj[a].push_back(b);
            Adj[b].push_back(a);
        }
        DFS(1,0,len - 1);
        Init_ST();
        while(m--){
            int a,b,c,k;sf("%d %d %d",&a,&b,&k);
            c = RMQ(a,b);
            pf("%d\n",Ar[Search(rts[a],rts[b],rts[c],rts[FA[c]],1,len - 1,k)]);
        }
    }
}