SPOJ COT - Count on a tree树链剖分加主席树

B - Count on a tree

Time Limit:129MS    Memory Limit:1572864KB    64bit IO Format:%lld & %llu

Submit Status Practice SPOJ COT

Description

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
  
    
    
    
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 

题目大意：给你一颗树，m次询问，问节点u到v第k大的数是多少。

思路：若是求一个数组的L到R的第k大的话，直接一个主席树就可以搞定了，现在变成了树上，然后其实是直接树剖，把u到v分成最多logn段的连续序列，对每个序列求出T[r] - T[l-1]然后相加起来，就是这logn段的左区间和，做法跟求数组的第k大一样。还有就是建主席树的时候放到树剖节点映射线段树那一部分，因为要保证重链上的节点是连续的。

#include <cstdio>
#include <cstring>
#include <queue>
#include <cmath>
#include <algorithm>
#include <iostream>
#include <cstdlib>
using namespace std;
//typedef __int64 LL;
const int maxn = 100005;
int w[maxn],fa[maxn],dep[maxn],num[maxn],son[maxn],top[maxn],vis[maxn],head[maxn],cnt,cou;
int T[maxn],lson[maxn*50],rson[maxn*50],c[maxn*50],tot,len,b[maxn],tt,user[maxn],usel[maxn];
int n,m;
struct Edge
{
    int v,next;
}e[maxn*2];
void add(int a,int b)
{
    e[cnt].v = b;
    e[cnt].next = head[a];
    head[a] = cnt++;
}
void init()
{
    cou = cnt = tot = len = 0;
    tt = 1;
    memset(head,-1,sizeof(head));
    memset(son,-1,sizeof(son));
}
void hase(int k)
{
    sort(b,b+k);
    len = unique(b,b+k) - b;
}
int get_hase(int now)
{
    return (lower_bound(b,b+len,now) - b);
}
int build(int l,int r)
{
    int root = tot++;
    c[root] = 0;
    if(l != r)
    {
        int mid = (l + r) >> 1;
        lson[root] = build(l,mid);
        rson[root] = build(mid+1,r);
    }
    return root;
}
void dfs1(int u,int v,int deep)
{
    fa[v] = u;
    num[v] = 1;
    dep[v] = deep;
    int i;
    for(i = head[v]; i != -1; i = e[i].next)
    {
        int end = e[i].v;
        if(end == u)
            continue;
        dfs1(v,end,deep+1);
        num[v] += num[end];
        if(son[v] == -1 || num[son[v]] < num[end])
        {
            son[v] = end;
        }
    }
}
int insert1(int root,int pos,int val)
{
    int newroot = tot++;
    c[newroot] = c[root] + val;
    int tmp = newroot;
    int l = 0,r = len - 1;
    while(l < r)
    {
        int mid = (l + r) >> 1;
        if(pos <= mid)
        {
            lson[newroot] = tot++;
            rson[newroot] = rson[root];
            root = lson[root];
            newroot = lson[newroot];
            r = mid;
        }
        else
        {
            lson[newroot] = lson[root];
            rson[newroot] = tot++;
            root = rson[root];
            newroot = rson[newroot];
            l = mid + 1;
        }
        c[newroot] = c[root] + val;
    }
    return tmp;
}
void dfs2(int u,int v)
{
    top[v] = u;
    T[tt] = insert1(T[tt-1],get_hase(w[v]),1);
    tt++;
    vis[v] = ++cou;
    if(son[v] != -1)
    {
        dfs2(u,son[v]);
    }
    else
        return;
    int i;
    for(i = head[v]; i != -1; i = e[i].next)
    {
        int end = e[i].v;
        if(end == son[v] || end == fa[v])
            continue;
        dfs2(end,end);
    }
}

void solve(int a,int b1,int pos)
{
    int q = top[a],w = top[b1];
    int t = 0;
    while(q != w)
    {
        if(dep[q] < dep[w])
        {
            swap(q,w);
            swap(a,b1);
        }
        user[t] = T[vis[a]],usel[t] = T[vis[q]-1];
        t++;
        a = fa[q];
        q = top[a];
    }
    if(dep[a] > dep[b1])
        swap(a,b1);
    user[t] = T[vis[b1]],usel[t] = T[vis[a]-1];
    t++;
    int l = 0,r = len - 1;
    int rootr,rootl,i;
    while(l < r)
    {
        int sum = 0;
        int mid = (l + r) >> 1;
        for(i = 0; i < t; i++)
        {
            sum += c[lson[user[i]]] - c[lson[usel[i]]];
        }
        if(sum < pos)
        {
            pos -= sum;
            l = mid + 1;
            for(i = 0; i < t; i++)
            {
                user[i] = rson[user[i]];
                usel[i] = rson[usel[i]];
            }
        }
        else
        {
            r = mid;
            for(i = 0; i < t; i++)
            {
                user[i] = lson[user[i]];
                usel[i] = lson[usel[i]];
            }
        }
    }
    printf("%d\n",b[l]);
}
int main()
{
    while(scanf("%d %d",&n,&m) != EOF)
    {
        int i;
        init();
        for(i = 1; i <= n; i++)
        {
            scanf("%d",&w[i]);
            b[len++] = w[i];

        }
        hase(len);
        int q,w;
        for(i = 1; i < n; i++)
        {
            scanf("%d %d",&q,&w);
            add(q,w);
            add(w,q);
        }
        T[0] = build(0,len-1);

        dfs1(1,1,1);
        dfs2(1,1);
        while(m--)
        {
            int A,B,C;
            scanf("%d %d %d",&A,&B,&C);
            solve(A,B,C);
        }
    }
    return 0;
}