백준 1717 Java 풀이

[1717-집합의 표현]

문제

• 입력받은 n에 대해 n + 1개의 서로소 집합이 존재한다. 0 ~ n 원소는 자기 자신
• 입력받은 m개의 커맨드가 입력된다. [연산] [타겟 숫자1] [타겟 숫자 2]
• 연산이 0일 경우 두 타겟 숫자가 속한 집합을 합집합 연산
• 연산이 1일 경우 두 타겟 숫자가 같은 집합인지 검사 후 YES or NO 출력

풀이

• Union Find를 사용하는 기본적인 문제다.
  

• 위 그림과 같은 두개의 집합은 각각 루트 노드 A, B를 부모 노드 값으로 갖고 있을것이다.

• 이 두개의 집합을 합집합 연산하기 위해선 B의 부모 노드를 A로 변경해주면 간단하다. find 연산시 부모 노드를 찾아 재귀를 돌기 때문에 타겟 노드 둘의 루트 노드를 찾으면 해결할 수 있게 되는 것이다.

• 핵심 메서드 union과 find다. nodes는 각 노드들의 배열

private void union(int a, int b) {
    int rootA = find(a);
    int rootB = find(b);
    
    //루드 노드를 찾아 루트를 병합
    if (rootA != rootB) {
        nodes[rootB] = rootA;
    }
}

private int find(int n) {
    if (nodes[n] == n) {
        return n;
    }
    //탐색 효율을 위해 n의 부모 노드를 전부 root노드로 변경
    return nodes[n] = find(nodes[n]);
}

코드

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.StringTokenizer;


public class Main {
    private static int[] nodes;

    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());

        nodes = new int[n + 1];

        for (int i = 0; i <= n; i++) {
            nodes[i] = i;
        }

        for (int i = 0; i < m; i++) {
            st = new StringTokenizer(br.readLine());
            String op = st.nextToken();
            int a = Integer.parseInt(st.nextToken());
            int b = Integer.parseInt(st.nextToken());

            if (op.equals("0")) {
                union(a, b);
            } else {
                if (find(a) == find(b)) {
                    System.out.println("YES");
                } else {
                    System.out.println("NO");
                }
            }
        }
    }

    private static void union(int a, int b) {
        int rootA = find(a);
        int rootB = find(b);

        if (rootA != rootB) {
            nodes[rootB] = rootA;
        }
    }

    private static int find(int n) {
        if (nodes[n] == n) {
            return n;
        }
        return nodes[n] = find(nodes[n]);
    }
}