SPOJ-COT Count on a tree(主席树 + LCA)

链接：https://www.spoj.com/problems/COT/en/

题意：给一颗树，每个点有对应的权值，每次查询形式为 u,v,k ，询问u到v的链上的第k小值是多少

分析：这道题就是一个主席树加上求u，v的LCA，主席树的建树，每个点以其父亲节点为基础建的，然后对于每次查询答案，先找到u,v的LCA，做一下差分，基本上是主席树的模板了。LCA是用树剖的方法求的，感觉比较好写，时间复杂度也可以接受。需要注意的是在求答案的时候，判断向左还是向右走与k值作比较，不是c[lson[ql]] + c[lson[qr]] - 2*c[lson[lca]]，也不是c[lson[ql]] + c[lson[qr]] - 2*c[lson[falca]]（falca为lca的父亲节点），应该是c[lson[ql]] + c[lson[qr]] - c[lson[lca]] - c[lson[falca]]，因为lca节点被算了两次但是只要减去一次，所以前两种方法都是错的。因为没有注意到这个问题，在这个地方WA了两次，仔细想了一下，才意识到这个问题。

代码：

#include <iostream>
#include <cstdio>
#include <cstring>
#include <algorithm>
#include <vector>
using namespace std;
int n,m;
int cnt,tot;
const int N = 100000 + 5;
vector<int>g[N];
int root[N],num[N],a[N],c[N*30],lson[N*30],rson[N*30];
int dep[N],fa[N],top[N],dfn[N],son[N],siz[N];
void init()
{
    for(int i = 1; i <= n; i++) num[i] = a[i];
    sort(a+1,a+1+n);
    cnt = unique(a+1,a+1+n) - (a + 1);
}
int build(int l,int r)
{
    int d = tot++;
    c[d] = 0;
    if(l == r) return d;
    int m = (l +r) >> 1;
    lson[d] = build(l,m);
    rson[d] = build(m+1,r);
    return d;
}
int insert(int d,int p)
{
    int nowd = tot ++ , res = nowd;
    int l = 1, r = cnt;
    c[nowd] = c[d] + 1;
    while(l < r)
    {
        int m = (l +r) >> 1;
        if(p <= m)
        {
            lson[nowd] = tot ++; rson[nowd] = rson[d];
            nowd = lson[nowd],d = lson[d];
            r = m;
        }
        else
        {
            rson[nowd] = tot ++; lson[nowd] = lson[d];
            nowd = rson[nowd]; d = rson[d];
            l = m + 1;
        }
        c[nowd] = c[d] + 1;
    }
    return res;
}
void dfs1(int u,int f,int d)
{
    dep[u] = d;
    fa[u] = f;
    son[u] = -1;
    siz[u] = 1;
    int p = lower_bound(a+1,a+1+cnt,num[u]) - a;
    root[u] = insert(root[f],p);
    for(int i = 0; i < g[u].size(); i++)
    {
        int v = g[u][i];
        if(v == f) continue;
        dfs1(v,u,d+1);
        siz[u] += siz[v];
        if(son[u] == -1 || siz[v] > siz[son[u]]) son[u] = v;
    }
    return;
}
void dfs2(int u,int t)
{
    top[u] = t;
    if(son[u] == -1) return;
    dfs2(son[u],t);
    for(int i = 0; i < g[u].size(); i++)
    {
        int v = g[u][i];
        if(v == fa[u] || v == son[u]) continue;
        dfs2(v,v);
    }
    return;
}
int getlca(int x,int y)
{
    int fx = top[x], fy = top[y];
    while(fx != fy)
    {
        if(dep[fx] < dep[fy]) swap(x,y) , swap(fx,fy);
        x = fa[fx]; fx = top[x];
    }
    if(dep[x] > dep[y]) swap(x,y);
    return x;
}
int query(int ql,int qr,int k)
{
    int lca = getlca(ql,qr);
    int lf = fa[lca];
    int l = 1, r = cnt;
    ql = root[ql], qr = root[qr], lca = root[lca],lf = root[lf];
    while(l < r)
    {
        int m = (l +r) >> 1;
        if(k <= c[lson[ql]] + c[lson[qr]] - c[lson[lca]] - c[lson[lf]])
        {
            ql = lson[ql], qr = lson[qr];lca = lson[lca]; lf = lson[lf];
            r = m;
        }
        else
        {
            k -= c[lson[ql]] + c[lson[qr]] - c[lson[lca]] - c[lson[lf]];
            ql = rson[ql], qr = rson[qr]; lca = rson[lca]; lf = rson[lf];
            l = m + 1;
        }
    }
    return l;
}
int main()
{
    while(~scanf("%d%d",&n,&m))
    {
        for(int i = 1; i <= n; i++) scanf("%d",&a[i]),g[i].clear();
        for(int i = 1; i < n; i++)
        {
            int x,y; scanf("%d%d",&x,&y);
            g[x].push_back(y);
            g[y].push_back(x);
        }
        init();
        tot = 0;
        root[0] = build(1,cnt);
        dfs1(1,0,0);
        dfs2(1,1);
        for(int i = 0; i < m; i++)
        {
            int x,y,k; scanf("%d%d%d",&x,&y,&k);
            int ans = query(x,y,k);
            printf("%d\n",a[ans]);
        }
    }
    return 0;
}