SPOJ - COT Count on a tree （主席树，树上第K小）

Count on a tree

You are given a tree with N nodes. The tree nodes are numbered from 1to N. Each node has an integer weight.

We will ask you to perform the following operation:

• u v k : ask for the kth minimum weight on the path from node u to node v

Input

In the first line there are two integers N and M. (N, M <= 100000)

In the second line there are N integers. The ith integer denotes the weight of the ith node.

In the next N-1 lines, each line contains two integers u v, which describes an edge (u, v).

In the next M lines, each line contains three integers u v k, which means an operation asking for the kth minimum weight on the path from node u to node v.

Output

For each operation, print its result.

Example

Input:
8 5
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 

Output:
2
8
9
105
7 

题意：给一棵树，每个点有一个权值，有q 个询问，每一次给定 u,v,k ，求u 到 v 路径上节点权值第 k 大的结点值。

思路：主席树模版题，在 dfs 求 LCA 的时候 update，最后判断的条件就是 c[lson[rr]] + c[lson[lr]] - c[lson[LCArt]] - c[lson[faLCArt]];

代码：

#include<cstdio>
#include<algorithm>
#include<cmath>
#include<queue>
#define debug(x) cout << "[" << #x <<": " << (x) <<"]"<< endl
#define pii pair<int,int>
#define clr(a,b) memset((a),b,sizeof(a))
#define rep(i,a,b) for(int i = a;i < b;i ++)
#define pb push_back
#define MP make_pair
#define LL long long
#define ull unsigned LL
#define ls i << 1
#define rs i << 1 + 1
#define INT(t) int t; scanf("%d",&t)

using namespace std;

const int maxn = 2e5 + 10;
const int M = maxn * 40;
int T[maxn];///T[i]表示第i颗线段树的顶结点
int lson[M],rson[M],c[M];///C[i]就代表第i个点的权值,即上述权值线段树的sum
int X[maxn],K[maxn];     /// lson[i], rson[i] 表示第 i 个点的左儿子,右儿子是谁
int cnt = 0,en;          /// cnt 用于动态开点

int head[maxn];
struct xx{
    int u,v,nex;
    xx(int u = 0,int v = 0,int nex = 0):
        u(u),v(v),nex(nex){}
}edge[maxn];
int coun = 0;
///倍增算法
int dep[maxn];            ///深度
int fa[maxn][25];         /// fa[i][j] 表示 i 节点的第 2^j 的祖先
int n,m,q;        ///其中 m = log(n * 1.0) / log(2.0);  n 为总数，编号为 1 ~ n;


int update(int root,int pos,int val){
    int newroot = cnt ++,tmp = newroot;
    c[newroot] = c[root] + val;
    int l = 1,r = en;
    while(l < r){
        int mid = (l + r) >> 1;
        if(pos <= mid){
            lson[newroot] = cnt ++;
            rson[newroot] = rson[root];
            newroot = lson[newroot];
            root = lson[root];
            r = mid;
        }

        else {
            rson[newroot] = cnt ++;
            lson[newroot] = lson[root];
            newroot = rson[newroot];
            root = rson[root];
            l = mid + 1;
        }
        c[newroot] = c[root] + val;
    }
    return tmp;
}

void dfs(int f,int s){
    dep[s] = dep[f] + 1;
    int pos = lower_bound(X + 1,X + 1 + en,K[s]) - X;
    T[s] = update(T[f],pos,1);
    for(int j = head[s];~j ;j = edge[j].nex){
        int v = edge[j].v;
        if(v != f){
            fa[v][0] = s;
            for(int i = 1;i <= m;i ++)
                fa[v][i] = fa[fa[v][i-1]][i-1];
            dfs(s,v);
        }
    }
}

int lca(int u,int v){
    if(dep[u] < dep[v]) swap(u,v);      ///u 在下面
    int d = dep[u] - dep[v];
    for(int i = 0; (1 << i) <= d;i ++)      ///移动到相同深度
        if((1 << i) & d)                ///(1<<i)&f找到f化为2进制后1的位置，移动到相应的位置
            u = fa[u][i];               ///比如f=5(101),先移动2^0祖先，然后再移动2^2祖先

    if(u != v){
        for(int i = m;i >= 0;i --)
            if(fa[u][i] != fa[v][i])         ///从最大祖先开始，判断a,b祖先，是否相同
                u = fa[u][i], v = fa[v][i];  ///如不相同，a b同时向上移动2^j
        u = fa[u][0];
    }
    return u;
}

int build(int l,int r){
    int root = cnt ++;
    c[root] = 0;
    if(l != r){
        int mid = (l + r) >> 1;
        lson[root] = build(l,mid);
        rson[root] = build(mid + 1,r);
    }
    return root;
}

int query(int u,int v,int k){
    int l = 1,r = en;
    int lroot = T[u];
    int rroot = T[v];
    int LCA = lca(u,v);
    int LCArt = T[LCA];
    int faLCArt = T[fa[LCA][0]];
    while(l < r){
        int mid = (l + r) >> 1;
        int tmp = c[lson[rroot]] + c[lson[lroot]] - c[lson[LCArt]] - c[lson[faLCArt]];
        if(tmp >= k){   ///简单画一画就理解了
            r = mid;
            lroot = lson[lroot];
            rroot = lson[rroot];
            LCArt = lson[LCArt];
            faLCArt = lson[faLCArt];
        }
        else {
            l = mid + 1;
            k -= tmp;
            rroot = rson[rroot];
            lroot = rson[lroot];
            LCArt = rson[LCArt];
            faLCArt = rson[faLCArt];
        }
    }
    return l;
}

int main() {
    scanf("%d%d",&n,&q);
    m = log(n*1.0) / log(2.0);
    for(int i = 1;i <= n;++ i){
        scanf("%d",&K[i]);
        X[i] = K[i];
    }
    sort(X + 1,X + 1 + n);
    en = unique(X + 1,X + 1 + n) - X - 1;
    for(int i = 0;i < maxn;++ i) head[i] = -1;
    for(int i = 1;i < n;++ i){
        int u,v; scanf("%d%d",&u,&v);
        edge[coun] = xx(u,v,head[u]);
        head[u] = coun ++;
        edge[coun] = xx(v,u,head[v]);
        head[v] = coun ++;
    }
    T[0] = build(1,en);
    dep[0] = -1;
    dfs(0,1);
    while(q --){
        int u,v,k; scanf("%d%d%d",&u,&v,&k);
        printf("%d\n",X[query(u,v,k)]);
    }
    return 0;
}