SPOJ COT Count on a tree（树上主席树 + LCA 求点第k小）题解

题意：n个点的树，每个点有权值，问你u~v路径第k小的点的权值是？

思路：

树上主席树就是每个点建一棵权值线段树，具体看JQ博客，LCA用倍增logn求出，具体原理看这里

树上主席树我每个点的存的是点u到源点1的权值线段树，那我求点u到v的所有点，显然是 u + v - lca - fa[lca]，就是u到1 + v到1 - 多算的lca - 多算的fa[lca]。不能减去两个lca不然少一个点了，

LCA板子：

//LCA
int fa[maxn][];
int dep[maxn];
void lca_dfs(int u, int pre, int d){
    dep[u] = d;
    fa[u][] = pre;
    for(int i = head[u]; i != -; i = edge[i].next)
        if(edge[i].v != pre)
            lca_dfs(edge[i].v, u, d + );
}
void lca_update(){
    for (int i = ; ( << i) <= n; i++)
        for(int u = ; u <= n; u++)
            fa[u][i] = fa[fa[u][i - ]][i - ];
}
int lca_query(int u, int v){
    if(dep[u] < dep[v]) swap(u, v);
    int d = dep[u] - dep[v];
    for(int i = ; ( << i) <= d; i++) {
        if(d & ( << i)) {
            u = fa[u][i];
        }
    }
    if(u != v) {
        for(int i = (int)log2(n); i >= ; i--) {
            if(fa[u][i] != fa[v][i]) {
                u = fa[u][i];
                v = fa[v][i];
            }
        }
        u = fa[u][];
    }
    return u;
}

代码：

#include<cmath>
#include<set>
#include<queue>
#include<cstdio>
#include<vector>
#include<cstring>
#include <iostream>
#include<algorithm>
using namespace std;
typedef long long ll;
typedef unsigned long long ull;
const int maxn =  + ;
const int M = maxn * ;
const ull seed = ;
const int INF = 0x3f3f3f3f;
const int MOD = ;
int n, m;
int root[maxn], a[maxn], tot;
struct Edge{
    int v, next;
}edge[maxn << ];
int head[maxn], tol;
void addEdge(int u, int v){
    edge[tol].v = v;
    edge[tol].next = head[u];
    head[u] = tol++;
}
struct node{
    int lson, rson;
    int sum;
}T[maxn * ];
void init(){
    memset(T, , sizeof(T));
    memset(root ,, sizeof(root));
    memset(head, -, sizeof(head));
    tot = tol = ;
}
vector<int> ve;
int getid(int x){
    return lower_bound(ve.begin(), ve.end(), x) - ve.begin() + ;
}
void update(int l, int r, int &now, int pre, int v, int pos){
    T[++tot] = T[pre], T[tot].sum += v, now = tot;
    if(l == r) return;
    int m = (l + r) >> ;
    if(pos <= m)
        update(l, m, T[now].lson, T[pre].lson, v, pos);
    else
        update(m + , r, T[now].rson, T[pre].rson, v, pos);
}
void build(int now, int pre){
    update(, n, root[now], root[pre], , getid(a[now]));
    for(int i = head[now]; i != -; i = edge[i].next){
        int v = edge[i].v;
        if(v == pre) continue;
        build(v, now);
    }
}
int query(int l, int r, int now, int pre, int lca, int flca, int k){
    if(l == r) return l;
    int m = (l + r) >> ;
    int sum = T[T[now].lson].sum + T[T[pre].lson].sum - T[T[lca].lson].sum - T[T[flca].lson].sum;
    if(sum >= k)
        return query(l, m, T[now].lson, T[pre].lson, T[lca].lson, T[flca].lson, k);
    else
        return query(m + , r, T[now].rson, T[pre].rson, T[lca].rson, T[flca].rson, k - sum);
}

//LCA
int fa[maxn][];
int dep[maxn];
void lca_dfs(int u, int pre, int d){
    dep[u] = d;
    fa[u][] = pre;
    for(int i = head[u]; i != -; i = edge[i].next)
        if(edge[i].v != pre)
            lca_dfs(edge[i].v, u, d + );
}
void lca_update(){
    for (int i = ; ( << i) <= n; i++)
        for(int u = ; u <= n; u++)
            fa[u][i] = fa[fa[u][i - ]][i - ];
}
int lca_query(int u, int v){
    if(dep[u] < dep[v]) swap(u, v);
    int d = dep[u] - dep[v];
    for(int i = ; ( << i) <= d; i++) {
        if(d & ( << i)) {
            u = fa[u][i];
        }
    }
    if(u != v) {
        for(int i = (int)log2(n); i >= ; i--) {
            if(fa[u][i] != fa[v][i]) {
                u = fa[u][i];
                v = fa[v][i];
            }
        }
        u = fa[u][];
    }
    return u;
}
int main(){
    init();
    ve.clear();
    scanf("%d%d", &n, &m);
    for(int i = ; i <= n; i++)
        scanf("%d", &a[i]), ve.push_back(a[i]);
    sort(ve.begin(), ve.end());
    ve.erase(unique(ve.begin(), ve.end()), ve.end());
    for(int i = ; i <= n - ; i++){
        int u, v;
        scanf("%d%d", &u, &v);
        addEdge(u, v);
        addEdge(v, u);
    }
    lca_dfs(, , );
    lca_update();
    build(, );
    while(m--){
        int u, v, k;
        scanf("%d%d%d", &u, &v, &k);
        int lca = lca_query(u, v);
        int ans = query(, n, root[u], root[v], root[lca], root[fa[lca][]], k);
        printf("%d\n", ve[ans - ]);
    }
    return ;
}