백준 1717 집합의 표현

문제

풀이

public class Q1717_집합 {

    static int[] arr;

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());
        int n = Integer.parseInt(st.nextToken());
        int m = Integer.parseInt(st.nextToken());
        arr = new int[n + 1];
        for (int i = 1; i <= n; n++) {
            arr[i] = i;
        }
        for (int i = 0; i < m; i++) {
            st = new StringTokenizer(br.readLine());
            // 유니온인지 파인드인지
            int type = Integer.parseInt(st.nextToken());
            int a = Integer.parseInt(st.nextToken());
            int b = Integer.parseInt(st.nextToken());
            // 유니온이라면
            if (type == 0) {
                union(a, b);
            // 파인드라면
            } else {
                if (check(a, b)) {
                    System.out.println("YES");
                } else {
                    System.out.println("NO");
                }
            }
        }
    }
	
    public static void union(int a, int b) {
    	// a의 집합대표와 b의 집합대표를 확인하고
        a = find(a);
        b = find(b);
        // 집합이 다르다면
        if (a != b) {
        	// a의 집합대표로 수정
            arr[b] = a;
        }
    }
    
    // 집합 대표를 찾는 연산
    // 재귀함수를 이용해 
    public static int find(int a) {
    	// 인덱스와 벨류가 같은것
        // 즉 집합 대표가 나오면 return
        if (a == arr[a]) {
            return a;
        // 집합 대표가 아니라면 재귀함수를 호출하고
        // 그 return값으로 탐색된 모든 인덱스의
        // 값(벨류)를 집합 대표값으로 변경
        } else {
            return arr[a] = find(arr[a]);
        }
    }
	
    // 두 집합이 같은지 확인하는 함수
    public static boolean check(int a, int b) {
        a = find(a);
        b = find(b);
        if (a == b) {
            return true;
        } else {
            return false;
        }
    }
}

리뷰

유니온 파인드(union find)라는 알고리즘을 이용해야 하는 문제였다.
유니온 각 노드가 속한 집합을 하나로 합치는 연산이며, 이때 합쳐진 값이 집합 대표가 된다.
파인드는 각 노드의 집합 대표를 반환하는 노드이다.

예를 들어 노드 a=1와 b=2일때 union(a,b)을 하면 a=1,b=1이 되고 이를 배열로 보면 arr = {1,1}이 된다.
이 때 find(b)를 하면 1이 반환된다.
여기서 집합 대표는 인덱스와 값이 동일한 노드를 말하는데, 위에서 들었던 예시에서 arr = {1,1}에서
배열이 1부터 시작한다고 했을 때 arr[1] = 1, arr[2] = 1이므로 arr[1]이 집합 대표가 되는 것이다.

이걸 적용한 풀이의 흐름은 다음과 같다.

1. 입력된 연산이 union일 때, 노드 a와 b의 집합 대표를 찾는다.

   • 이때 집합 대표를 구하기 위해서는, find연산이 필요하므로 find연산을 호출한다.

2. find 연산은 노드의 집합 대표를 반환한다.

   • find연산의 자세한 역할은 아래에서 확인하자.

3. 입력된 연산이find일 때, 노드 a와 b의 집합 대표가 같다면 YES를 다르다면 NO를 출력한다.

• 여기서 헷갈렸던 부분이 find연산이었다. find연산은 2가지 역할을 수행한다.

1. 집합 대표 찾기
2. 집합 대표와 같은 집합인 노드의 모든 값을 집합 대표의 값으로 수정하기
   • 만약 노드가 집합 대표라면 해당 값을 반환한다.
   • 다르다면 집합 대표가 나올 때까지 자기 자신 호출한다.
   • 집합 대표를 찾았다면 재귀호출을 거슬러 올라가며 재귀호출을 한 모든 노드의 벨류를 집합 대표로 바꿔준다.
     예를 들어, 배열 arr = {1,2,3}이 있을 때 union(2,3)을 수행하면 arr = {1,2,2}가 되고,
     union(1,2)를 수행하면 arr = {1,1,2}가 된다.
     이 상태에서 find(3)을 수행한다면 반환값은 2가 아니라 1이 되어야 한다.
     따라서 arr[3] = 1로 수정하여 arr = {1,1,1}이, 반환값은 1이 된다.