Count on a tree 树上区间第K小

Count on a tree

题意：求路径 u到v上的 第k小的权重。

题解：先DFS建数， 然后对于每个节点往上跑出一颗主席树， 然后每次更新。 查询的时候， u， v， k， 找到  z = lca(u,v) , p = anc[z][0], 然后对于这条路上左边子节点的个数就可以被表示为u，v的树- z,p的树上的值。

然后主要是LCA不怎么用写搓了， 找了好久的bug。 汗。

代码：

 #include<bits/stdc++.h>
 using namespace std;
 #define Fopen freopen("_in.txt","r",stdin); freopen("_out.txt","w",stdout);
 #define LL long long
 #define ULL unsigned LL
 #define fi first
 #define se second
 #define pb push_back
 #define max3(a,b,c) max(a,max(b,c))
 #define min3(a,b,c) min(a,min(b,c))
 typedef pair<int,int> pll;
 const int INF = 0x3f3f3f3f;
 const LL mod =  (int)1e9+;
 const int N = 1e5 + ;
 const int M = 5e6 + ;
 vector<int> son[N];
 int lson[M], rson[M], cnt[M];
 int anc[N][];
 int tot, t;
 int a[N], w[N], deep[N], root[N];
 int id(int x){
     return lower_bound(a+, a++t, x) - a;
 }
 int Build(int l, int r){
     int now = ++tot;
     cnt[now] = ;
     if(l < r){
         int m = l+r >> ;
         lson[now] = Build(l, m);
         rson[now] = Build(m+, r);
     }
     return now;
 }
 int Update(int l, int r, int pre, int c, int v){
     int now = ++tot;
     cnt[now] = cnt[pre] + v;
     if(l < r){
         int m = l+r >> ;
         if(c <= m){
             rson[now] = rson[pre];
             lson[now] = Update(l, m, lson[pre], c, v);
         }
         else {
             lson[now] = lson[pre];
             rson[now] = Update(m+, r, rson[pre], c, v);
         }
     }
     return now;
 }

 void dfs(int o, int u){
     deep[u] = deep[o] + ;
     root[u] = Update(, t, root[o], id(w[u]), );
     for(int i = ; i < son[u].size(); i++){
         int v = son[u][i];
         if(v == o) continue;
         anc[v][] = u;
         for(int j = ; j < ; j++) anc[v][j] = anc[anc[v][j-]][j-];
         dfs(u, v);
     }
 }
 int lca(int u, int v){
     if(deep[u] < deep[v]) swap(u, v);
     for(int i = ; i >= ; i--)
         if(deep[anc[u][i]] >= deep[v]) u = anc[u][i];
     if(u == v) return v;
     for(int i = ; i >= ; i--)
         if(anc[u][i] != anc[v][i])
             u = anc[u][i], v = anc[v][i];
     return anc[u][];
 }
 int Query(int l, int r, int c1, int c2, int d1, int d2, int k){
     if(l == r) return a[l];
     int m = l+r >> ;
     int num = cnt[lson[c1]] + cnt[lson[c2]] - cnt[lson[d1]] - cnt[lson[d2]];
     if(num >= k) return Query(l, m, lson[c1], lson[c2], lson[d1], lson[d2], k);
     else return Query(m+, r, rson[c1], rson[c2], rson[d1], rson[d2], k-num);
 }
 int main(){
     int n, q, u, v, k, p, z;
     scanf("%d%d", &n, &q);
     for(int i = ; i <= n; i++){
         scanf("%d", &w[i]);
         a[i] = w[i];
     }
     for(int i = ; i < n; i++){
         scanf("%d%d", &u, &v);
         son[u].pb(v);
         son[v].pb(u);
     }
     sort(a+, a++n);
     for(int i = ; i <= n; i++) a[++t] = a[i];
     root[] = Build(, t);
     dfs(, );
     while(q--){
         scanf("%d%d%d", &u, &v, &k);
         z = lca(u, v);
         p = anc[z][];
         printf("%d\n", Query(, t, root[u], root[v], root[z], root[p], k));
     }
     return ;
 }

Count on a tree