SP10628 COT - Count on a tree

Posted 2021-03-13 sxy_limit

主席树的综合运用题.

前置芝士

1. 可持久化线段树:其实就是主席树了.
2. LCA:最近公共祖先,本题需要在(log_2N)及以内的时间复杂度内解决这个问题.

具体做法

主席树维护每个点到根节点这一条链上不同树出现的次数,然后发现这个东西是可以相减的,于是这条链上每个数出现的次数就变成了(sum[u]+sum[v]-2*sum[LCA(u,v)]).然后就可以发现这个是错的,如果按这个式子计算最后的链上就没有LCA位置的值了,所以在范围包含(val[LCA(u,v)])时需要加一,然后正常去找kth就好了.

代码

#include<bits/stdc++.h>
#define REP(i,first,last) for(int i=first;i<=last;++i)
#define DOW(i,first,last) for(int i=first;i>=last;--i)
using namespace std;
const int maxN=2e5+7;

int N,M,point[maxN];
int root[maxN];
int sor[maxN];
map<int,int>Hash;
int val[maxN];
//一下为主席树部分
int point_cnt=0;
struct Tree//对于树的一个结构体
{
    int sum,lson,rson;
}tree[maxN*32];
//主席树标准define
#define LSON tree[now].lson
#define RSON tree[now].rson
#define MIDDLE ((left+right)>>1)
#define LEFT LSON,left,MIDDLE
#define RIGHT RSON,MIDDLE+1,right
#define NEW_LSON tree[new_tree].lson
#define NEW_RSON tree[new_tree].rson
void PushUp(int now)//合并信息
{
    tree[now].sum=tree[LSON].sum+tree[RSON].sum;
}
void UpData(int num,int &new_tree,int now,int left=1,int right=N)
//修改并产生一个新版本到new_tree中
{
    if(num<left||right<num)
    {
        new_tree=now;
        return;
    }
    new_tree=++point_cnt;
    if(left==right)
    {
        tree[new_tree].sum=tree[now].sum+1;
        return;
    }
    UpData(num,NEW_LSON,LEFT);
    UpData(num,NEW_RSON,RIGHT);
    PushUp(new_tree);
}
int QueryKth(int k,int LCA_num,int tree1/*相对于式子中的u*/,int tree2/*相当于式子中的v*/,int tree3/*相当于式子中的LCA(u,v)*/,int left=1,int right=N)
{
    if(left==right)
    {
        return left;
    }
    int sum=tree[tree[tree1].lson].sum
           +tree[tree[tree2].lson].sum
         -2*tree[tree[tree3].lson].sum;//计算出这条链上在left~right中的数出现的次数
    if(left<=LCA_num&&LCA_num<=MIDDLE)//如果LCA位置的数在当前的范围内就加一
    {
        sum++;
    }
    if(sum>=k)//左子树查询kth
    {
        return
        QueryKth(k,LCA_num,
        tree[tree1].lson,
        tree[tree2].lson,
        tree[tree3].lson,
        left,MIDDLE);
    }
    return//右子树查询
    QueryKth(k-sum,LCA_num,
    tree[tree1].rson,
    tree[tree2].rson,
    tree[tree3].rson,
    MIDDLE+1,right);
}
//一下为链式前向星部分(非主要内容,不多讲)
int head[maxN];
int edge_cnt=0;
struct Edge
{
    int to,next;
}edge[maxN];
#define FOR(now) for(int i_edge=head[now];i_edge;i_edge=edge[i_edge].next)
#define TO edge[i_edge].to
void AddEdge(int form,int to)
{
    edge[++edge_cnt].to=to;
    edge[edge_cnt].next=head[form];
    head[form]=edge_cnt;
}
//一下为倍增LCA部分(非主要内容,不多讲)
int kth_father[maxN][25];
bool visit[maxN];
int deep[maxN],father[maxN];
void DFS(int now)
{
    UpData(Hash[point[now]],root[now],root[father[now]]);//在DFS的过程中把主席树建好
    deep[now]=deep[father[now]]+1;
    kth_father[now][0]=father[now];
    REP(i,0,22)
    {
        kth_father[now][i+1]=kth_father[kth_father[now][i]][i];
    }
    FOR(now)
    {
        if(father[now]!=TO)
        {
            father[TO]=now;
            DFS(TO);
        }
    }
}
int LCA(int x,int y)//LCA核心代码
{
    if(deep[x]<deep[y])
    {
        swap(x,y);
    }
    DOW(i,22,0)
    {
        if(deep[kth_father[x][i]]>=deep[y])
        {
            x=kth_father[x][i];
        }
        if(x==y)
        {
            return x;
        }
    }
    DOW(i,22,0)
    {
        if(kth_father[x][i]!=kth_father[y][i])
        {
            x=kth_father[x][i];
            y=kth_father[y][i];
        }
    }
    return father[x];
}

int main()
{
    scanf("%d%d",&N,&M);
    REP(i,1,N)
    {
        scanf("%d",&point[i]);
        sor[i]=point[i];
    }
    int fa,son;
    REP(i,1,N-1)
    {
        scanf("%d%d",&fa,&son);
        AddEdge(fa,son);
        AddEdge(son,fa);
    }
    sort(sor+1,sor+1+N);//离散化
    sor[0]=-114514;
    int cnt_num=0;
    REP(i,1,N)
    {
        if(sor[i]!=sor[i-1])
        {
            Hash[sor[i]]=++cnt_num;
            val[cnt_num]=sor[i];
        }
    }
    N=cnt_num;
    DFS(1);//DFS,建树+倍增LCA预处理
    int u,v,k;
    int point_LCA;
    REP(i,1,M)
    {
        scanf("%d%d%d",&u,&v,&k);
        point_LCA=LCA(u,v);//找到LCA
        printf("%d
",
            val[QueryKth(k,Hash[point[point_LCA]],
                     root[u],root[v],root[point_LCA])]
            );//直接查找kth
    }
    return 0;
}