[Codeforces 1006E] Military Problem

[Codeforces 1006E] Military Problem - 1600

문제 설명 : 

군대의 위계를 나타내는 트리가 주어지고, 쿼리가 주어진다. 각 쿼리는 두 개의 정수 u,k로 이루어져 있는데 쿼리마다 노드 u부터 시작해서 k번째로 방문하는 노드가 몇 번인지를 찾는 문제이다. 방문하는 원칙에는 반드시 현재 노드보다 위계가 낮은(아래)노드로만 방문이 가능하며, 자식이 여럿일 때는 노드 번호가 작은 순부터 방문한다.

풀이 : 

트리의 노드가 최대 20만개, 쿼리도 최대 20만개이므로 매 쿼리를 계산할 때 최대 O(logN)이내에 수행해야 한다. 따라서 전처리 단계부터 쿼리에 빠르게 답할 수 있도록 처리해야 한다. 그러기 위해서 나는 두 개의 배열을 마련할 것이다. 하나는 (밑의 소스코드와 같이) cnt[i]=a라는 배열로 루트노드부터 시작해서 i번째 방문하는 노드의 번호는 a라는 의미이고, node_cnt[j]=b라는 배열은 이와 반대로 j번째 노드를 방문하기 위해서는 루트노드부터 b번만큼 움직여야 한다는 의미이다. 또한 각 노드는 자신의 자식 노드가 몇 개인지를 저장하는 sz[]배열도 갖는다.  이 배열을 dfs로 전처리를 돌면서 모두 채워준다. 다 채웠으면 매 쿼리를 받을 것이다. 우선 각 쿼리에서 u를 받았으면 그 노드 u까지 도달하는데  걸리는 횟수 (= node_cnt[u] )를 구하고, 그 횟수에 +(k-1)을 해주면 그 결과값을 cnt[]배열에 인덱스에 넣으면 그 의미는 루트노드에서부터 시작하여 그 계산한 횟수만큼 뒤에 몇 번 노드에 도착하는 지를 알 수 있게 되고, 이게 곧 정답이다. k에 1을 빼주는 이유는 dfs를 돌 때 맨 처음 시작은 자기 자신부터 시작하기 때문이다.

주의할 점 : 

-

배운 점 : 

루트노드로부터 도달 횟수를 관리하는 배열 cnt[]와, 이와 반대로 해당 노드까지 도달하는데 걸리는 횟수를 관리하는 node_cnt[]를 사용하여 풀 수 있었다.

코드 :

#include<stdio.h>
#include<math.h>
#include<string.h>
#include<iostream>
#include<functional>
#include<string>
#include<algorithm>
#include<vector>
#include<map>
#include<set>
#include<queue>
#include<assert.h>
#include<stdlib.h>
#include<stack>
using namespace std;
/*********************Contest Template***********************/
typedef long long ll;
typedef unsigned long long ull;
typedef pair<int,int> pii;
typedef pair<long long,long long> pll;
typedef pair<int,int> Cord;
#define FASTIO ios_base::sync_with_stdio(false);cin.tie(0);
struct S{
    int a,b; S(){}S(int _a,int _b){
        a=_a;
        b=_b;
    }
    const bool operator<(const S &o) const{
        return a<o.a;
    }
};
priority_queue<int,vector<int>,greater<int>> mpq;
#define SPACE 0
#define NL 1
inline void showAll(vector<int> &v,int NL_or_space){ 
    //0 is space, 1 is NL
    for(int &here:v){
        printf("%d",here);
        printf("%c",(NL_or_space)?'\n':' ');
    }
}
/*********************Contest Template***********************/
const int SIZE= 200009;
 
vector<int> tree[SIZE];
int sz[SIZE],node_cnt[SIZE],cnt[SIZE];
int total=0;
 
int dfs(int par,int node){
    node_cnt[node]=++total;
    cnt[total]=node;
 
    sort(tree[node].begin(),tree[node].end());
 
    for(int next:tree[node]){
        if(next!=par)
            sz[node]+=dfs(node,next);
    }
 
    return ++sz[node];
}
 
int main(){
    int n,q;    scanf("%d %d",&n,&q);
 
    for(int i=2 ; i<=n ; i++){
        int inp;    scanf("%d",&inp);
        tree[inp].push_back(i);
        tree[i].push_back(inp);
    }
    dfs(0,1);
 
    while(q--){
        int u,k;    scanf("%d %d",&u,&k);
 
        if(sz[u]<k) printf("-1\n");
        else{
            printf("%d\n",cnt[node_cnt[u]+k-1]);
        }
    }
    return 0;
}