SPOJ - COT Count on a tree （主席树）

题目链接

You are given a tree with N nodes. The tree nodes are numbered from 1 to N. Each node has an integer weight.

We will ask you to perform the following operation:

• u v k : ask for the kth minimum weight on the path from node u to node v

Input

In the first line there are two integers N and M. (N, M <= 100000)

In the second line there are N integers. The ith integer denotes the weight of the ith node.

In the next N-1 lines, each line contains two integers u v, which describes an edge (u, v).

In the next M lines, each line contains three integers u v k, which means an operation asking for the kth minimum weight on the path from node u to node v.

Output

For each operation, print its result.

Example

Input:
8 5
105 2 9 3 8 5 7 7
1 2
1 3
1 4
3 5
3 6
3 7
4 8
2 5 1
2 5 2
2 5 3
2 5 4
7 8 2 

Output:
2
8
9
105
7 

题意：题意很好理解，就是求一个树上两节点路径中，权值第k小的数。

题解：很多博客都是直接给出一个公式但是没有解释是为什么（可能这就是大佬），下图大概就是这棵树的结构。首先我们要求两个节点之间路径的第k小，肯定会涉及到两个节点的LCA。比如求（8,5）两节点中间的数，路径为8->1->5。这样我们最基本的思路就出来了，我们从根节点开始建树，然后，儿子节点来继承父节点。这样建树建完了过后，就可以直接查询了。因为主席树有点类似于前缀和（自我理解），如果我们要求[l,r]的和，是不是应该sum[r]-sum[l-1];。这里情况也差不多，我们用root数组来记录主席树根节点，我们求1-8之间有多少点就是root[8]-root[father[1]]（这里father数组存的是每个节点的父节点）。因为1节点我们在这里算了，所以下面我们只算3-5之间的数，他们之间的数就是root[5]-root[1]。这样算下来5-8路径之间的数为root[8]-root[father[1]]+root[5]-root[1];后面就是常规的主席树了。

所以我们可以得出公式（u,v）之间的数字个数为：

root[u]+root[v]-root[LCA]-root[[father[LCA]];

 

#include <iostream>
#include<cstring>
#include<cstdio>
#include<algorithm>
#include<map>
#include<queue>
#include<set>
#include<cmath>
#include<stack>
#include<string>
const int maxn=1e5+5;
const int mod=10007;
const int inf=1e9;
const long long onf=1e18;
#define me(a,b) memset(a,b,sizeof(a))
#define lowbit(x) x&(-x)
#define lson l,mid,rt<<1
#define rson mid+1,r,rt<<1|1
#define PI 3.14159265358979323846
typedef long long ll;
typedef unsigned long long ull;
using namespace std;
int n,m,cnt,tot,len;
int a[maxn],ta[maxn];
int father[maxn][30],depth[maxn];
int root[maxn];
int head[maxn];
struct node{
    int u,v,next;
    node(int _u=0,int _v=0,int _next=0):u(_u),v(_v),next(_next){}
}maps[maxn<<1];
struct Node{
    int sum,ls,rs;
}tree[40*maxn];
void init(){
    me(head,-1);
    cnt=tot=0;
}
void add_edge(int u,int v){
    maps[cnt].v=v;
    maps[cnt].next=head[u];
    head[u]=cnt++;
}
void push_date(int pre,int &node,int pos,int l,int r){///跟新
    node=++tot;
    tree[node].ls=tree[pre].ls,tree[node].rs=tree[pre].rs;
    tree[node].sum=tree[pre].sum+1;
    if(l==r)
        return ;
    int mid=(l+r)>>1;
    if(pos<=mid)
        push_date(tree[pre].ls,tree[node].ls,pos,l,mid);
    else
        push_date(tree[pre].rs,tree[node].rs,pos,mid+1,r);
}
int query(int x,int y,int fa_lca,int lca,int k,int l,int r){
    if(l==r)
        return ta[l];
    int temp=tree[tree[x].ls].sum+tree[tree[y].ls].sum-tree[tree[fa_lca].ls].sum-tree[tree[lca].ls].sum;
    int mid=(l+r)>>1;
    if(k<=temp)
        return query(tree[x].ls,tree[y].ls,tree[fa_lca].ls,tree[lca].ls,k,l,mid);
    else
        return query(tree[x].rs,tree[y].rs,tree[fa_lca].rs,tree[lca].rs,k-temp,mid+1,r);
}
void init_dfs(int u,int fa){///DFS建树和对深度和father数组预处理。
    depth[u]=depth[fa]+1;
    father[u][0]=fa;
    for(int i=1;(1<<i)<=n;i++)
        father[u][i]=father[father[u][i-1]][i-1];
    int pos=lower_bound(ta+1,ta+len+1,a[u])-ta;///离散化
    push_date(root[fa],root[u],pos,1,len);
    for(int i=head[u];i!=-1;i=maps[i].next){
        int v=maps[i].v;
        if(v==fa)
            continue ;
        init_dfs(v,u);
    }
}
int LCA(int u,int v){///求LCA
    if(depth[u]>depth[v])
        swap(u,v);
    int ret=depth[v]-depth[u];
    for(int i=0;(1<<i)<=n;i++){
        if(ret&(1<<i))
            v=father[v][i];
    }
    if(u==v)
        return v;
    for(int i=log2(n);i>=0;i--){
        if(father[u][i]==father[v][i])
            continue ;
        u=father[u][i],v=father[v][i];
    }
    return father[u][0];
}
int main()
{
    init();
    scanf("%d%d",&n,&m);
    for(int i=1;i<=n;i++){
        scanf("%d",&a[i]);
        ta[i]=a[i];
    }
    sort(ta+1,ta+n+1);
    len=unique(ta+1,ta+n+1)-ta-1;
    for(int i=1;i<n;i++){
        int u,v;
        scanf("%d%d",&u,&v);
        add_edge(u,v),add_edge(v,u);
    }
    init_dfs(1,0);
    while(m--){
        int u,v,k;
        scanf("%d%d%d",&u,&v,&k);
        int temp_father=LCA(u,v);///求两点的LCA
        printf("%d\n",query(root[u],root[v],root[father[temp_father][0]],root[temp_father],k,1,len));
    }
    return 0;
}