BOJ 22996 - 유니온 파인드 복원 링크
(2024.03.21 기준 G2)

문제

$1$부터 $N$까지의 원소를 하나씩 가지는 $N$개의 집합 $\{1\}$, $\{2\}$, $\cdots$, $\{N\}$이 있다. 그리고 아래와 같이 주어지는 질의를 순서대로 $Q$개 처리해야 한다.

• $1$번 질의: $1$ $u$ $v$ 의 형식으로 주어진다. 원소 $u$가 속한 집합과 원소 $v$가 속한 집합이 다르다면, 두 집합을 하나로 합친다.
• $2$번 질의: $2$ $v$ 의 형식으로 주어진다. 원소 $v$가 속한 집합에 있는 원소를 아무거나 하나 반환한다.

위와 같은 입력 데이터를 처리했을 때의 $par$ 배열과 $N$, $Q$, $2$번 질의들에 대한 결과만 알 때, 주어진 입력 데이터를 복원해서 출력

알고리즘

분리 집합의 성질과 위상 정렬을 이용한 해 구성 문제

풀이

$2$번 질의는 나중에 섞이면 복잡해지니깐 처음부터 바로 처리하자. 초기 상태에는 집합마다 하나의 원소만 있으니깐 $2$번 질의들이 반환한 값들을 다시 질의 형태로 출력하면 된다.

문제에서의 merge 함수를 보자. merge(int u, int v)이면 $u$가 속한 집합이 $v$가 속한 집합에 속하게 된다. 즉, $par(u) \rightarrow par(u)$ 유향 간선이라고 생각할 수 있다.

이제 주어지는 $par$ 배열에서 $i \ne par(i)$인 모든 $i$와 $par(i)$에 대해 유향 간선 $i \rightarrow par(i)$가 필요함을 알 수 있다.

하지만 여기서 문제가 있다. 만약 $par$ 배열이 $[1, 1, 2]$이라고 생각해보자. 그러면 여기서 필요한 유향 간선은 $2 \rightarrow 1$과 $3 \rightarrow 2$이다. 여기서 서술한 순서대로 질의를 출력하면 어떻게 될까?

• $2 \rightarrow 1$을 처리하면 $par$ 배열은 $[1, 1, 3]$
• $3 \rightarrow 2$를 처리하면 $par$ 배열은 $[1, 1, 1]$. $2$는 이미 $1$을 가리키고 있기 때문이다.
  이는 결국 진입차수가 $0$이 아닌 정점을 먼저 처리해서 일어나는 문제다.

위와 같은 문제 때문에 결국은 유향 그래프의 유향 간선들을 위상 정렬 순서에 따라 간선을 출력해야 한다.

이제 위 두 과정을 거치면 하나의 문제가 남는다. 바로 현재 출력한 질의의 개수가 $Q$보다 부족한 경우.
이는 $par(v) = v$인 $v$를 하나 찾고, $1$ $v$ $v$를 부족한 개수만큼 출력하면 된다.

일단 $par$ 배열에서의 유향 간선들은 곧 포레스트 형태를 이룬다. 즉 무조건 $par(v) = v$인 $v$가 하나 이상 존재한다. 조건을 만족하는 $v$는 자신이 속한 그래프(트리)에서 루트이다. 여기서 찾은 $v$를 자기 자신한테로 merge를 시도하면 $par(v)$는 바뀌지 않을 뿐더러 merge 함수도 실행되지 않게 된다.

코드

• C++

#include <bits/stdc++.h>
using namespace std;

int main(){
    ios_base::sync_with_stdio(0);
    cin.tie(0);

    int N, Q; cin >> N >> Q; cout << N << ' ' << Q << '\n';
    int par[N + 1]; for (int i = 1; i <= N; i++) cin >> par[i];

    // 2번 질의는 처음에 바로 처리하자.
    int M; cin >> M; // M은 2번 질의의 개수이자, 현재 출력한 질의의 개수로 사용하자.
    for (int _ = 0, v; _ < M; _++){
        cin >> v;
        cout << 2 << ' ' << v << '\n';
    }

    /* par 배열은 곧 i->par[i] 유향 간선으로 생각할 수 있다.
    u->v->w 순서인데 v->w, u->v로 처리하면 u->w, v->w가 형태가 되어버린다.
    즉, 순서에 맞게 간선을 출력해야 한다.
    유향 간선은 위상 정렬로 처리할 수 있다. */

    vector<int> G[N + 1]; // 유향 그래프
    int ind[N + 1]; fill(ind + 1, ind + N + 1, 0); // 정점마다 진입차수
    for (int i = 1; i <= N; i++) if (par[i] != i){
        ++M; G[i].push_back(par[i]); ++ind[par[i]]; // 필요한 유향 간선을 저장
    }

    stack<int> stk; // 진입차수가 현재 0인 정점들을 담을 스택
    for (int i = 1; i <= N; i++) if (!ind[i]) stk.push(i);

    while (!stk.empty()){
        // u는 진입차수가 0인 정점이라서 이제 u에서 나아가는 간선을 출력해도 된다.
        int u = stk.top(); stk.pop();
        for (int v: G[u]){
            cout << 1 << ' ' << u << ' ' << v << '\n';
            if (!--ind[v]) stk.push(v); // v의 진입차수가 0이 되었다면 스택에 담는다.
        }
    }

    // 현재 출력한 질의의 개수가 Q보다 부족하다면
    // 그만큼 v와 par[v]가 같은 v를 출력하면 된다.
    if (M < Q){
        int v;
        for (int i = 1; i <= N; i++) if (par[i] == i){
            v = i;
            break;
        }
        for (; M < Q; M++) cout << 1 << ' ' << v << ' ' << v << '\n';
    }
}

• Python

import sys; input = sys.stdin.readline

N, Q = map(int, input().split()); print(N, Q)
par = [0] + list(map(int, input().split()))

# 2번 질의는 처음에 바로 처리하자.
M = int(input()) # M은 2번 질의의 개수이자, 현재 출력한 질의의 개수로 사용하자.
for _ in range(M):
    v = int(input())
    print(2, v)

''' par 배열은 곧 i->par[i] 유향 간선으로 생각할 수 있다.
u->v->w 순서인데 v->w, u->v로 처리하면 u->w, v->w가 형태가 되어버린다.
즉, 순서에 맞게 간선을 출력해야 한다.
유향 간선은 위상 정렬로 처리할 수 있다. '''

G = [[] for _ in range(N + 1)] # 유향 그래프
ind = [0] * (N + 1) # 정점마다 진입차수
for i in range(1, N + 1):
    if par[i] != i:
        M += 1; G[i].append(par[i]); ind[par[i]] += 1 # 필요한 유향 간선을 저장

stk = [] # 진입차수가 현재 0인 정점들을 담을 스택
for i in range(1, N + 1):
    if not ind[i]:
        stk.append(i)

while stk:
    # u는 진입차수가 0인 정점이라서 이제 u에서 나아가는 간선을 출력해도 된다.
    u = stk.pop()
    for v in G[u]:
        print(1, u, v)
        ind[v] -= 1
        if not ind[v]:
            stk.append(v) # v의 진입차수가 0이 되었다면 스택에 담는다.

# 현재 출력한 질의의 개수가 Q보다 부족하다면
# 그만큼 v와 par[v]가 같은 v를 출력하면 된다.
if M < Q:
    for i in range(1, N + 1):
        if par[i] == i:
            v = i
            break
    for _ in range(M, Q):
        print(1, v, v)