SPOJ - COT Count on a tree LCA+主席树 树上路径求第K大模板题

SPOJ - COT Count on a tree

题意：

给出一棵树，每个点都自己的权重，然后给出树上的边，要求从节点 u 到节点 v 路径上的第 k 小的权重的大小。
因为权重可能很大，所以需要离散化

思路：

在树上建主席树

每个点建一棵权值线段树，记录从这个点到根的路径上的点权信息。u可以和fa[u]的线段树共用很多信息，所以上主席树！那么 u 到 v 路径上的点权信息，就是tree[u]+tree[v]-tree[lc]-tree[fa[lc][0]]，lc 表示 lca(u,v)。
 

#include<iostream>
#include<algorithm>
#include<cstdlib>
#include<sstream>
#include<cstring>
#include<bitset>
#include<cstdio>
#include<string>
#include<deque>
#include<stack>
#include<cmath>
#include<queue>
#include<set>
#include<map>
#define mod 1000000007
using namespace std;
typedef long long ll;
const int maxn = 1e5+10;
struct edge
{
    int to;
    int next;
}e[2*maxn];
struct node
{
    int l,r;
    int sum;
}t[40*maxn];
int n,m,len,a[maxn],b[maxn],rt[maxn],tot; // 主席树
int cnt,head[maxn],fa[maxn][20],d[maxn],N;  //LCA
// ***** 主席树 *****
void build(int x,int l,int r,int &p)
{
    t[++tot] = t[p];
    p = tot;
    t[p].sum++;
    if(l==r)
        return ;
    int mid = (l+r)>>1;
    if(x<=mid)
        build(x,l,mid,t[p].l);
    else
        build(x,mid+1,r,t[p].r);
}
int query(int p1,int p2,int p3,int p4,int k,int l,int r)
{
    if(l==r)
        return l;
    int val = t[t[p1].l].sum+t[t[p2].l].sum-t[t[p3].l].sum-t[t[p4].l].sum;
    int mid = (l+r)>>1;
    if(k<=val)
        return query(t[p1].l,t[p2].l,t[p3].l,t[p4].l,k,l,mid);
    else
        return query(t[p1].r,t[p2].r,t[p3].r,t[p4].r,k-val,mid+1,r);
}
// ***** LCA *****
void add(int u,int v)
{
    e[cnt].to = v;
    e[cnt].next = head[u];
    head[u] = cnt++;
}
void dfs(int u,int f)
{
    for(int i=1;i<=N;i++)
        fa[u][i] = fa[fa[u][i-1]][i-1];
    rt[u] = rt[fa[u][0]];
    build(a[u],1,len,rt[u]);
    for(int i=head[u];i!=-1;i=e[i].next)
    {
        int v = e[i].to;
        if(v==f)
            continue;
        d[v] = d[u]+1;  //深度
        fa[v][0] = u;
        dfs(v,u);
    }
}
int lca(int u,int v)
{
    if(d[u]<d[v])
        swap(u,v);
    int dc = d[u]-d[v];
    for(int i=0;i<=N;i++)
    {
        if( (1<<i)&dc )
            u = fa[u][i];
    }
    if(u==v)
        return u;
    for(int i=N;i>=0;i--)
    {
        if(fa[u][i]!=fa[v][i])
        {
            u = fa[u][i];
            v = fa[v][i];
        }
    }
    return fa[u][0];
}

int main()
{
    int u,v;
    int k,lc;
    cnt = tot = 0;
    memset(head,-1,sizeof(head));
    scanf("%d%d",&n,&m);
    N = (int)(log(1.0*n)/log(2))+1;
    for(int i=1;i<=n;i++)
    {
        scanf("%d",&a[i]);
        b[i] = a[i];
    }
    sort(b+1,b+n+1);
    len = unique(b+1,b+n+1)-b-1;   //离散化
    for(int i=1;i<=n;i++)
        a[i] = lower_bound(b+1,b+len+1,a[i])-b;
    for(int i=1;i<n;i++)
    {
        scanf("%d%d",&u,&v);
        add(u,v);
        add(v,u);
    }
    dfs(1,0);
    for(int i=1;i<=m;i++)
    {
        scanf("%d%d%d",&u,&v,&k);
        lc = lca(u,v);
        printf("%d\n",b[query(rt[u],rt[v],rt[lc],rt[fa[lc][0]],k,1,len)]);
    }
    return 0;
}