최대 원소의 개수가 1,000,000이고 질의 개수가 100,000으로 상당히 큰 편이다. 따라서 경로 압축이 필요한 전형적인 유니온 파인드 문제이다.
유니온 파인드란 여러 노드가 있을 때 특정 2개의 노드를 연결해 1개의 집합으로 묶는 union 연산과 두 노드가 같은 집합에 속해 있는지 확인하는 find 연산으로 구성되어 있는 알고리즘이다.
유니온 파인드를 표현하는 일반적인 방법은 1차원 배열을 이용하는 것이다. 처음에는 모든 노드가 연결되어 있지 않으므로 각 노드가 대표 노드가 된다. 따라서 아래와 같이 배열은 자신의 인덱스 값으로 초기화한다.
| 1 | 2 | 3 | 4 | 5 | 6 |
|---|---|---|---|---|---|
| 1 | 2 | 3 | 4 | 5 | 6 |
1, 4에 대해 union 연산을 진행한다고 생각해보자. 1은 대표 노드, 4는 자식 노드로 union 연산을 진행하게 되면 배열 [4]의 대표 노드를 1로 설정하게 된다.
동일하게 5, 6에 대해 union 연산을 진행한다고 생각해보자. 5는 대표 노드, 6은 자식 노드로 union 연산을 진행하게 되면 배열 [6]의 대표노드를 5로 설정하게 된다.
그러면 집합을 표현한 1차원 배열은 아래와 같이 표현된다.
| 1 | 2 | 3 | 4 | 5 | 6 |
|---|---|---|---|---|---|
| 1 | 2 | 3 | 1 | 5 | 5 |
이제 4, 6에 대해 union 연산을 진행한다고 생각해보자. 그런데 4, 6은 대표노드가 아니다. 따라서 각 노드의 대표 노드를 찾아 올라간 다음, 대표 노드끼리 연결을 진행해야 한다.
이와 같이 자신이 속한 집합의 대표노드를 찾는 연산을 find라고 한다.
따라서 4의 대표 노드 1에 6의 대표 노드 5를 연결하여 아래와 같은 배열이 된다.
| 1 | 2 | 3 | 4 | 5 | 6 |
|---|---|---|---|---|---|
| 1 | 2 | 3 | 1 | 1 | 5 |
find 연산은 아까 설명했듯이 자신이 속한 집합의 대표 노드를 찾는 연산이다. 그런데 단순히 대표 노드를 찾는 역할만 하는 것이 아니라, 그래프를 정돈하고 시간복잡도를 향상시키는 기능도 있다.
다음은 find 연산의 작동 원리이다.
- 대상 노드 배열에 index값과 value값이 동일한지 확인한다.
- 동일하지 않으면 value값이 가리키는 index 위치로 이동한다.
- 이동 위치의 index값과 value값이 같을 때까지 1. ~ 2.를 반복한다. 반복이므로 해당 부분은 재귀 함수로 구현하게 된다.
- 대표 노드에 도달하면 재귀 함수를 빠져나오면서 거치는 모든 노드 값을 로트 노드 값으로 변경한다.
대표 노드에 도달하면 재귀 함수를 나오면서 거치는 모든 노드는 루트 노드 값으로 변경하게 되므로 6번 노드 역시 대표 노드가 1이 된다.
| 1 | 2 | 3 | 4 | 5 | 6 |
|---|---|---|---|---|---|
| 1 | 2 | 3 | 1 | 1 | 1 |
결국 여러번 find를 진행하게 되면 그래프에서 경로 압축이 발생하여 시간 복잡도가 O(1)로 줄어드는 효과를 보인다.
알고리즘에 대해 설명한다고 설명이 좀 길었는데, 이제 문제를 풀어보도록 하자.
예제 입력을 기준으로 문제를 풀어보겠다. 초기 집합의 개수는 0 ~ 7까지 주어져 있으므로 초기 배열은 다음과 같다.
| 0 | 1 | 2 | 3 | 4 | 5 | 6 | 7 |
|---|---|---|---|---|---|---|---|
| 0 | 1 | 2 | 3 | 4 | 5 | 6 | 7 |
이후 총 8번의 연산이 이루어지는데, 0 a b은 집합을 합치는 연산이므로 우리가 앞서 배운 union을 진행하면 되고, 1 a b의 경우 같은 집합인지 확인하는 연산이므로 find 연산을 통해 a와 b의 대표 노드가 서로 같으면 YES, 다르면 NO를 출력하면 된다.
따라서 전체 소스코드는 다음과 같다.
import java.util.Scanner;
public class P_1717 {
static int[] parent;
public static void main(String[] args) {
Scanner sc = new Scanner(System.in);
int N = sc.nextInt();
int M = sc.nextInt();
parent = new int[N + 1];
for (int i = 0; i <= N; i++) {
parent[i] = i;
}
for (int i = 0; i < M; i++) {
int op = sc.nextInt();
int a = sc.nextInt();
int b = sc.nextInt();
if (op == 0)
union(a, b);
else {
if (checkSame(a, b))
System.out.println("YES");
else
System.out.println("NO");
}
}
}
static void union(int a, int b) {
a = find(a);
b = find(b);
if (a != b) {
parent[a] = b;
}
}
static int find(int a) {
if (parent[a] == a)
return a;
else
return parent[a] = find(parent[a]);
}
static boolean checkSame(int a, int b) {
a = find(a);
b = find(b);
if (a == b)
return true;
return false;
}
}코드에서 find 함수를 보면 주어진 노드 a가 대표 노드와 같을 경우 a 자체가 대표 노드이므로 그대로 a를 반환하면 된다.
만약 다를 경우 자식 노드이므로 a의 대표 노드에 대해 find 연산을 재귀적으로 진행하여 찾은 값을 자식 노드 a의 대표 노드로 설정해 주면 된다.
코드는 간결한 편이지만 원리를 확실히 파악해야 하니 잘 기억하도록 하자.
