[BOJ] 13306 트리

문제: https://www.acmicpc.net/problem/13306

 2016 KOI 중등부 3번 문제인 트리이다. 고등부 3번과 굉장히 비슷한 문제이기는 한데, 듣기로는 고등부 3번은 HLD라는 알고리즘을 써야 한다고 알고 있는데 이 문제는 상대적으로 깔끔하게 발상의 전환을 통해 문제를 해결할 수 있다. 실제로 대회장에 갔을 때도 모든 풀이는 다 생각해 놨는데 자료구조 구현을 못해서 실패했다. 웬만하면 풀이를 보기 전에 생각을 더 해봤으면 하는 좋은 문제.

 내 해법은 Union-Find Tree로 알려져 있는 자료구조 (Disjoint Set)을 구현해서 문제를 역순으로 풀어 나가는 것이다. 어떤 정점과 그 정점의 부모 정점을 제거하면 그 두 트리는 분리되게 된다. 두 정점 사이에 경로가 존재하려면 결과적으로 그 둘의 최소 공통 조상을 지나야 하므로, 그 두 정점은 같은 서브트리에 속해야 한다는 말이다. 결과적으로 경로가 있는지 확인하는 작업을 두 정점의 루트가 같은지 다른지 비교하는 것만으로도 할 수 있다.

 그런데 중요한 것은 이 문제는 정점들을 서로 분리하는 것이므로 Union-Find Tree로 루트 정점을 빠르게 찾아보려고 해도 Collapsing Find를 사용할 수가 없다... 경로압축이 되면 결과적으로는 분리되던 말던 맨 위 루트 정점을 반환하게 되는데 이러면 어떤 트리를 임의의 두 서브트리로 나누었을때 바뀌는 루트를 빠르게 표현해줄 수가 없다. 

 여기에서 발상의 전환이 필요하다. 만약에 어떤 트리를 두 서브트리로 쪼개면 각 서브트리의 모든 정점은 반대쪽 서브트리의 정점들과는 완전히 연결이 끊기게 된다. 핵심은 이때 정점 개수-1만큼 반드시 쪼개는 연산이 실행되므로 최종적인 결과는 모든 트리가 루트만을 원소로 가지는 트리가 된다는 것이다. 그럼 이제 이러한 발상으로 이어갈 수 있다. 맨 처음에는 모든 정점을 다 따로따로 떨어트려 놓고, 쿼리를 역순으로 읽으면서 거꾸로 간선을 붙여가면서 연결되는지 확인하면 되지 않을까? 이러면 트리를 분리가 아닌 합치는 연산이 되므로 Union-Find Tree를 효율적으로 사용할 수 있다. 이 코드는 다음과 같은 생각을 직접 옮긴 코드이다. 난이도는 3/5정도 주고 싶다.

#include <cstdio>
#include <vector>
#include <algorithm>
 
using namespace std;
 
typedef struct DisJointSet
{
    vector<int> parent;
    vector<int> ra;
    DisJointSet(int n)
    {
        parent.assign(n+1,0);
        ra.assign(n+1,0);
        for(int i=0;i<=n;i++)
            parent[i]=i;
    }
    int fi(int a)
    {
        if(parent[a]==a) return a;
        else return parent[a]=fi(parent[a]);
    }
    void un(int a, int b)
    {
        a=fi(a); b=fi(b);
        if(a==b) return;
        if(ra[a]>ra[b]) swap(a,b);
        parent[b]=a;
        if(ra[a]==ra[b]) ra[b]++;
    }
}DJS;
 
int n,m;
int query[400005][3];
 
int main()
{
    scanf("%d %d",&n,&m);
    DJS S(n);
    vector<int> parent(n+1);
    vector<bool> ans;
    parent[1]=1;
    for(int i=2;i<=n;i++)
        scanf("%d",&parent[i]);
    for(int i=0;i<n+m-1;i++)
    {
        scanf("%d",&query[i][0]);
        if(query[i][0]==0)
            scanf("%d",&query[i][1]);
        else
            scanf("%d %d",&query[i][1],&query[i][2]);
    }
    for(int i=n+m-2;i>=0;i--)
    {
        int t=query[i][0];
        if(t==0)
        {
            int t1=query[i][1];
            S.un(t1,parent[t1]);
        }
        else
        {
            int t1=query[i][1],t2=query[i][2];
            if(S.fi(t1)==S.fi(t2))
                ans.push_back(true);
            else
                ans.push_back(false);
        }
    }
    reverse(ans.begin(),ans.end());
    for(int i=0;i<ans.size();i++)
    {
        if(ans[i]) printf("YES\n");
        else printf("NO\n");
    }
    return 0;
}