[BZOJ 2588][SPOJ COT]Count On a Tree(DFS序主席树)

题目链接

http://www.lydsy.com/JudgeOnline/problem.php?id=2588

思路

静态DFS序主席树模板题。
就是对整个树求出它的DFS序，然后把这个DFS序看成一个普通的序列，然后对于这个DFS序的每个前缀各建立一个线段树(对于某些i

代码

#include <iostream>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <algorithm>

#define MAXN 100010
#define MAXM 2100000

using namespace std;

struct edge
{
    int u,v,next;
}edges[MAXN*2];

int head[MAXN],nCount1=0;

void AddEdge(int U,int V)
{
    edges[++nCount1].u=U;
    edges[nCount1].v=V;
    edges[nCount1].next=head[U];
    head[U]=nCount1;
}

int lc[MAXM],rc[MAXM],sum[MAXM],root[MAXN];
int stack[MAXN],top=0,hash[MAXN],tot=0;

int findhash(int x)
{
    int lowerBound=1,upperBound=tot;
    while(lowerBound<=upperBound)
    {
        int mid=(lowerBound+upperBound)>>1;
        if(hash[mid]==x) return mid;
        else if(hash[mid]>x) upperBound=mid-1;
        else lowerBound=mid+1;
    }
    return lowerBound;
}

int fa[MAXN][20];
int depth[MAXN],num[MAXN],pos[MAXN],len=0; //num数组保存DFS序，pos[i]=点i在DFS序中的位置

void dfs(int u) //DFS将无根树转有根树
{
    num[++len]=u;
    pos[u]=len;
    for(int p=head[u];p!=-1;p=edges[p].next)
    {
        int v=edges[p].v;
        if(fa[u][0]==v) continue;
        depth[v]=depth[u]+1;
        fa[v][0]=u;
        dfs(v);
    }
}

int n,m;

void LCA_prework() //!!!!!
{
    for(int i=1;i<=16;i++)
        for(int j=1;j<=n;j++)
            fa[j][i]=fa[fa[j][i-1]][i-1];
}

int LCA(int a,int b) //求a和b的LCA
{
    if(depth[a]<depth[b]) swap(a,b);
    for(int i=16;i>=0;i--)
    {
        if(depth[fa[a][i]]<=depth[b]) continue;
        a=fa[a][i];
    }
    if(depth[a]!=depth[b]) a=fa[a][0];
    if(a==b) return a;
    for(int i=16;i>=0;i--)
    {
        if(!fa[a][i]||!fa[b][i]) continue;
        if(fa[a][i]==fa[b][i]) continue;
        a=fa[a][i],b=fa[b][i];
    }
    a=fa[a][0],b=fa[b][0];
    return a;
}

int nCount=0;

void update(int fa,int &o,int L,int R,int x,int val) //将x处的数字修改为val
{
    o=++nCount;
    sum[o]=sum[fa]+val;
    if(L==R) return;
    lc[o]=lc[fa],rc[o]=rc[fa];
    int M=(L+R)>>1;
    if(x<=M) update(lc[fa],lc[o],L,M,x,val);
    else update(rc[fa],rc[o],M+1,R,x,val);
}

int query(int x,int y,int k) //求x到y路径上的第k大
{
    int a=x,b=y,c=LCA(a,b),d=fa[c][0];
    a=root[pos[a]],b=root[pos[b]],c=root[pos[c]],d=root[pos[d]];
    int L=1,R=tot;
    while(L<R)
    {
        int M=(L+R)>>1;
        int tmp=sum[lc[a]]+sum[lc[b]]-sum[lc[c]]-sum[lc[d]];
        if(tmp>=k) R=M,a=lc[a],b=lc[b],c=lc[c],d=lc[d];
        else k-=tmp,L=M+1,a=rc[a],b=rc[b],c=rc[c],d=rc[d];
    }
    return hash[L];
}

int val[MAXN];

int main()
{
    memset(head,-1,sizeof(head));
    scanf("%d%d",&n,&m);
    for(int i=1;i<=n;i++)
    {
        scanf("%d",&val[i]);
        stack[i]=val[i];
    }
    sort(stack+1,stack+n+1);
    hash[++tot]=stack[1];
    for(int i=2;i<=n;i++)
        if(stack[i]!=stack[i-1])
            hash[++tot]=stack[i];
    for(int i=1;i<=n;i++) val[i]=findhash(val[i]); //离散化每个点的点权
    for(int i=1;i<n;i++)
    {
        int u,v;
        scanf("%d%d",&u,&v);
        AddEdge(u,v);
        AddEdge(v,u);
    }
    dfs(1);
    LCA_prework(); //!!!!
    for(int i=1;i<=n;i++)
    {
        int t=num[i];
        update(root[pos[fa[t][0]]],root[i],1,tot,val[t],1); //!!!!!!
    }
    int last=0;
    for(int i=1;i<=m;i++)
    {
        int x,y,K;
        scanf("%d%d%d",&x,&y,&K);
        x^=last;
        last=query(x,y,K);
        printf("%d",last);
        if(i!=m) printf("\n");
    }
    return 0;
}