[백준] 1717 - 집합의 표현 (Java)

문제 파악

https://www.acmicpc.net/problem/1717

풀이

집합을 합치는 행위를 수행하거나, 두 원소가 같은 집합인지 리턴해주기만 하면 된다

루트 개념을 가져가면 쉽게 풀이할 수 있다.
말로 풀어 설명하자면 부분집합의 대표(루트)를 설정하고, union 시 하나의 대표로 병합해 준다.
이전에 풀이한 최소 스패닝 트리 문제와 유사하게 풀이했다. (https://diggingcode.tistory.com/70)

이런 경우 대체로 1차원 배열을 선언해서 풀이하는 편이다.

초기에는 각 원소들이 부분집합의 대표이므로 배열에 스스로의 인덱스를 할당해준다.
이를 탐색 및 union 시 업데이트 해 주도록 하자.

초기 배열 상태와 union 시 업데이트 되는 배열의 상태

재귀적으로 루트를 찾는 함수를 호출하도록 구현하였다.

코드

import java.io.*;

public class Main {
    private static int[] nums;

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        String[] input = br.readLine().split(" ");
        int n = Integer.parseInt(input[0]);
        int m = Integer.parseInt(input[1]);
        nums = new int[n + 1];
        for (int i = 0; i <= n; i++) {
            nums[i] = i;
        }
        BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
        for (int i = 0; i < m; i++) {
            String[] line = br.readLine().split(" ");
            int operation = Integer.parseInt(line[0]);
            int x = Integer.parseInt(line[1]);
            int y = Integer.parseInt(line[2]);
            int a = Math.min(x, y);
            int b = Math.max(x, y);
            if (operation == 1) {
                bw.append(find(a) == find(b) ? "YES" : "NO").append("\n");
            } else {
                union(a, b);
            }
        }
        bw.flush();
    }

    static int find(int i) {
        if (nums[i] == i) {
            return i;
        }
        return nums[i] = find(nums[i]);
    }

    static void union(int i, int j) {
        int iRoot = find(i);
        int jRoot = find(j);
        if (iRoot == jRoot) return;
        nums[iRoot] = jRoot;
    }
}