牛客__樹的顏色

鏈接：https://ac.nowcoder.com/acm/problem/16547
來源：牛客網
 

時間限制：C/C++ 4秒，其他語言8秒
空間限制：C/C++ 262144K，其他語言524288K
64bit IO Format: %lld

題目描述

給定一棵樹，1爲根，樹上每個點有一個顏色

 

給q次詢問，每次詢問以x爲根的子樹中深度在[dep[x],dep[x]+y]之間的點的顏色集合(相同顏色只算一遍)中編號第k小的顏色是什麼，無解輸出-1

輸入描述:

第一行,兩個數n q
第二行n-1個數，表示2到n的父親
第三行n個數，表示n個點的顏色
接下來q行
每行三個數x y k，表示一次詢問

輸出描述:

輸出q行，對於每個詢問，輸出一個數表示答案

示例1

輸入

複製

5 3
1 1 2 2
1 2 1 3 2
1 1 2
1 2 3
2 1 1

輸出

複製

2
3
2

備註:

n<=105,q<=105

顏色在[1,n]之間

ps：整體二分法，將顏色分成塊，用樹狀數組維護每一塊的數目，遍歷塊數就可以得出答案，同時對於當前節點先加入非最大子樹的點（這樣就可以先算出非最大子樹的答案），然後再將最大子樹加入，因爲加入非最大子樹的點在查完後都要清掉，留下最大子樹就可以少刪很多點，看代碼理解吧

#include <bits/stdc++.h>
 
using namespace std;
 
typedef pair<int,int>P;
 
const int maxn = 1e5+9;
const int INF = 0x3f3f3f3f;
 
struct rt{
    int v,next;
}edge[maxn];
 
struct node{
    int y,k,id;
};
int n,q;
int tot,S;
int head[maxn],col[maxn],in[maxn],mm[maxn];
int bit[500][maxn];
int dep[maxn],siz[maxn],son[maxn];
 
void add_edge(int u,int v){
    edge[tot].v=v;
    edge[tot].next=head[u];
    head[u]=tot++;
}
 
vector<node>Q[maxn];
int ans[maxn];
 
void dfn(int u,int fa){
    dep[u]=dep[fa]+1;
    siz[u]=1;
    for(int i=head[u];~i;i=edge[i].next){
        int v=edge[i].v;
        dfn(v,u);
        siz[u]+=siz[v];
        if(siz[son[u]]<siz[v])son[u]=v;
    }
}
 
void update(int *p,int l,int x){
    while(l<=n){
        p[l]+=x;
        l+=(l&-l);
    }
}
 
int query(int *p,int x,int y){
    int sum=0;
    x--;
    while(x>=1){
        sum-=p[x];
        x-=(x&-x);
    }
    while(y>=1){
       sum+=p[y];
       y-=(y&-y);
    }
    return sum;
}
 
void upd(int u){
    if(mm[col[u]]<=dep[u])return;
    if(mm[col[u]]!=INF)update(bit[in[col[u]]],mm[col[u]],-1);
    mm[col[u]]=dep[u];
    update(bit[in[col[u]]],dep[u],1);
}
 
void get(int u){
    for(int i=head[u];~i;i=edge[i].next)get(edge[i].v);upd(u);
}
 
void del(int u){
    for(int i=head[u];~i;i=edge[i].next)del(edge[i].v);
    if(mm[col[u]]==INF)return;
    if(mm[col[u]]!=INF)update(bit[in[col[u]]],mm[col[u]],-1);
    mm[col[u]]=INF;
}
 
void dfs(int u,int fp=1){
    for(int i=head[u];~i;i=edge[i].next){
        int v=edge[i].v;
        if(v!=son[u])dfs(v,0);
    }
    if(son[u])dfs(son[u]);
    for(int i=head[u];~i;i=edge[i].next){
        int v=edge[i].v;
        if(v!=son[u])get(v);
    }
    upd(u);
    for(int i=0;i<Q[u].size();i++){
        int x=dep[u],y=dep[u]+Q[u][i].y,k=Q[u][i].k;
        y=min(y,n);
        ans[Q[u][i].id]=-1;
        for(int j=0;j<S;j++){
            int sum=query(bit[j],x,y);
            if(k>sum){k-=sum;continue;}
            for(int l=j*S;;l++){
                int w=(mm[l]>=x&&mm[l]<=y);
                if(w<k){k-=w;continue;}
                ans[Q[u][i].id]=l;break;
            }
            break;
        }
    }
    if(fp)return; del(u);
}
 
int main(){
    memset(head,-1,sizeof(head));
    memset(mm,0x3f,sizeof(mm));
    scanf("%d%d",&n,&q);
    int u;
    for(int i=2;i<=n;i++){
        scanf("%d",&u);
        add_edge(u,i);
    }
    S=(int)sqrt(n)+1;
    for(int i=1;i<=n;i++){
        scanf("%d",&col[i]);
        in[i]=i/S;
    }
    dep[1]=1;
    dfn(1,0);
    int x,y,k;
    for(int i=0;i<q;i++){
        scanf("%d%d%d",&x,&y,&k);
        Q[x].push_back(node{y,k,i});
    }
    dfs(1);
    for(int i=0;i<q;i++){
        printf("%d\n",ans[i]);
    }
    return 0;
}