【路径第k小 && 主席树】SPOJ - COT Count on a tree

Step1 Problem:

给你一颗 n 个节点的树，每个节点有相应的权值。
有 m 个询问，每次询问，问你 u 到 v 路径上第 k 小的权值。

Step2 Ideas:

核心思想：dist[i] 代表根节点到 i 的距离, LCA(u, v) 代表 u 和 v 的最近公共祖先。
如果我们要求 u 到 v 的距离 = dist[u] + dist[v] - 2 * dist[LCA(u, v)];
主席树：
每一颗线段树都是在父亲线段树的基础上修改logn个节点，区间维护的是权值（也就是权值线段树）。
这时我们要求 u 到 v 路径上第 k 小的权值，我们只需要弄出 u 到 v 路径对应的权值线段树即可。
也就是 root[u] + root[v] - root[LCA(u, v)] - root[LCA(u, v)的父亲]

Step3 Code:

#include<bits/stdc++.h>
using namespace std;
#define MID int mid = (l+r)/2
const int N = 1e5+5;
const int LN = 20;
struct node
{
    int sum, l, r;
}tree[N*40];
vector<int> Map[N];
int a[N], n, sca[N], ns, dep[N], fq[N][20];
int root[N], top = 0;
int creat(int sum, int l, int r)
{
    tree[top++] = (node){sum, l, r};
    return top-1;
}
void updata(int &root, int pre, int l, int r, int pos, int v)
{
    root = creat(tree[pre].sum+v, tree[pre].l, tree[pre].r);
    if(l == r) return ;
    MID;
    if(pos <= mid) updata(tree[root].l, tree[pre].l, l, mid, pos, v);
    else updata(tree[root].r, tree[pre].r, mid+1, r, pos, v);
}
int query(int root1, int root2, int root3, int root4, int l, int r, int k)
{
    if(l == r) return l;
    int l1 = tree[root1].l, l2 = tree[root2].l, l3 = tree[root3].l, l4 = tree[root4].l;
    int tmp = tree[l1].sum + tree[l2].sum - tree[l3].sum - tree[l4].sum;
    MID;
    if(tmp >= k) return query(l1, l2, l3, l4, l, mid, k);
    else return query(tree[root1].r, tree[root2].r, tree[root3].r, tree[root4].r, mid+1, r, k-tmp);

}
int get_id(int x)
{
    return lower_bound(sca+1, sca+1+ns, x) - sca;
}
void dfs(int u, int f, int d)
{
    fq[u][0] = f; dep[u] = d;
    updata(root[u], root[f], 1, ns, get_id(a[u]), 1);//在父亲线段树的基础上建线段树
    for(int i = 1; i < LN; i++)
        fq[u][i] = fq[fq[u][i-1]][i-1];
    for(int i = 0; i < Map[u].size(); i++)
    {
        int to = Map[u][i];
        if(to != f) dfs(to, u, d+1);
    }
}
int lca(int x, int y)
{
    if(dep[x] > dep[y]) swap(x, y);
    for(int i = 0; i < LN; i++)
    {
        if((dep[y]-dep[x])>>i&1) {
            y = fq[y][i];
        }
    }
    if(x == y) return x;
    for(int i = LN-1; i >= 0; i--)
    {
        if(fq[x][i] != fq[y][i]) {
            x = fq[x][i];
            y = fq[y][i];
        }
    }
    return fq[x][0];
}
int main()
{
    int m, u, v;
    scanf("%d %d", &n, &m);
    for(int i = 1; i <= n; i++)
    {
        scanf("%d", &a[i]);
        sca[i] = a[i];
    }
    sort(sca+1, sca+1+n);
    ns = unique(sca+1, sca+1+n) - sca - 1;
    for(int i = 1; i < n; i++)
    {
        scanf("%d %d", &u, &v);
        Map[u].push_back(v);
        Map[v].push_back(u);
    }
    root[0] = creat(0, 0, 0);
    dfs(1, 0, 0);
    int k;
    while(m--)
    {
        scanf("%d %d %d", &u, &v, &k);
        int Lca = lca(u, v);
        printf("%d\n", sca[query(root[u], root[v], root[Lca], root[fq[Lca][0]], 1, ns, k)]);//核心
    }
    return 0;
}