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

Posted 2021-10-21 kirinsb

题意：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][20];
int dep[maxn];
void lca_dfs(int u, int pre, int d){
    dep[u] = d;
    fa[u][0] = pre;
    for(int i = head[u]; i != -1; i = edge[i].next)
        if(edge[i].v != pre)
            lca_dfs(edge[i].v, u, d + 1);
}
void lca_update(){
    for (int i = 1; (1 << i) <= n; i++)
        for(int u = 1; u <= n; u++)
            fa[u][i] = fa[fa[u][i - 1]][i - 1];
}
int lca_query(int u, int v){
    if(dep[u] < dep[v]) swap(u, v);
    int d = dep[u] - dep[v];
    for(int i = 0; (1 << i) <= d; i++) {
        if(d & (1 << i)) {
            u = fa[u][i];
        }
    }
    if(u != v) {
        for(int i = (int)log2(n); i >= 0; i--) {
            if(fa[u][i] != fa[v][i]) {
                u = fa[u][i];
                v = fa[v][i];
            }
        }
        u = fa[u][0];
    }
    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 = 100000 + 10;
const int M = maxn * 30;
const ull seed = 131;
const int INF = 0x3f3f3f3f;
const int MOD = 1000000007;
int n, m;
int root[maxn], a[maxn], tot;
struct Edge{
    int v, next;
}edge[maxn << 1];
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 * 40];
void init(){
    memset(T, 0, sizeof(T));
    memset(root ,0, sizeof(root));
    memset(head, -1, sizeof(head));
    tot = tol = 0;
}
vector<int> ve;
int getid(int x){
    return lower_bound(ve.begin(), ve.end(), x) - ve.begin() + 1;
}
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) >> 1;
    if(pos <= m)
        update(l, m, T[now].lson, T[pre].lson, v, pos);
    else
        update(m + 1, r, T[now].rson, T[pre].rson, v, pos);
}
void build(int now, int pre){
    update(1, n, root[now], root[pre], 1, getid(a[now]));
    for(int i = head[now]; i != -1; 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) >> 1;
    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 + 1, r, T[now].rson, T[pre].rson, T[lca].rson, T[flca].rson, k - sum);
}

//LCA
int fa[maxn][20];
int dep[maxn];
void lca_dfs(int u, int pre, int d){
    dep[u] = d;
    fa[u][0] = pre;
    for(int i = head[u]; i != -1; i = edge[i].next)
        if(edge[i].v != pre)
            lca_dfs(edge[i].v, u, d + 1);
}
void lca_update(){
    for (int i = 1; (1 << i) <= n; i++)
        for(int u = 1; u <= n; u++)
            fa[u][i] = fa[fa[u][i - 1]][i - 1];
}
int lca_query(int u, int v){
    if(dep[u] < dep[v]) swap(u, v);
    int d = dep[u] - dep[v];
    for(int i = 0; (1 << i) <= d; i++) {
        if(d & (1 << i)) {
            u = fa[u][i];
        }
    }
    if(u != v) {
        for(int i = (int)log2(n); i >= 0; i--) {
            if(fa[u][i] != fa[v][i]) {
                u = fa[u][i];
                v = fa[v][i];
            }
        }
        u = fa[u][0];
    }
    return u;
}
int main(){
    init();
    ve.clear();
    scanf("%d%d", &n, &m);
    for(int i = 1; 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 = 1; i <= n - 1; i++){
        int u, v;
        scanf("%d%d", &u, &v);
        addEdge(u, v);
        addEdge(v, u);
    }
    lca_dfs(1, 0, 1);
    lca_update();
    build(1, 0);
    while(m--){
        int u, v, k;
        scanf("%d%d%d", &u, &v, &k);
        int lca = lca_query(u, v);
        int ans = query(1, n, root[u], root[v], root[lca], root[fa[lca][0]], k);
        printf("%d\\n", ve[ans - 1]);
    }
    return 0;
}