[백준 P1717] 집합의 표현 Java

P1717 집합의 표현

1717번: 집합의 표현

문제

초기에 {0}, {1}, {2}, ... {n} 이 각각 n+1개의 집합을 이루고 있다. 여기에 합집합 연산과, 두 원소가 같은 집합에 포함되어 있는지를 확인하는 연산을 수행하려고 한다.

집합을 표현하는 프로그램을 작성하시오.

입력값

첫째 줄에 n(1 ≤ n ≤ 1,000,000), m(1 ≤ m ≤ 100,000)이 주어진다. m은 입력으로 주어지는 연산의 개수이다. 다음 m개의 줄에는 각각의 연산이 주어진다. 합집합은 0 a b의 형태로 입력이 주어진다. 이는 a가 포함되어 있는 집합과, b가 포함되어 있는 집합을 합친다는 의미이다. 두 원소가 같은 집합에 포함되어 있는지를 확인하는 연산은 1 a b의 형태로 입력이 주어진다. 이는 a와 b가 같은 집합에 포함되어 있는지를 확인하는 연산이다. a와 b는 n 이하의 자연수 또는 0이며 같을 수도 있다.

7 8
0 1 3
1 1 7
0 7 6
1 7 1
0 3 7
0 4 2
0 1 1
1 1 1

출력값

1로 시작하는 입력에 대해서 한 줄에 하나씩 YES/NO로 결과를 출력한다. (yes/no 를 출력해도 된다)

NO
NO
YES

알고리즘

{0}, ….. {n} 까지 n + 1개의 집합이 존재한다. 그리고 두 개의 집합을 합하는 연산과 두 원소가 같은 집합에 포함되어 있는지를 확인하는 연산을 구현해야 한다.

전형적인 Union-Find 문제이다. Union으로 두 개의 집합을 합하는 연산을 구현하고 Find로 두 원소가 같은 집합에 포함되어 있는지를 확인하면 된다.

우선 n + 1개의 원소가 몇 번째 집합에 속해있는지 저장할 배열을 선언한다. 배열은 parent로 이름을 정했다.

parent = new int[N + 1];

for (int i = 1; i <= N; i++) {
    parent[i] = i;
}

find 연산을 구현한다. find는 매개변수로 받아오는 원소가 속한 집합의 대표값을 반환한다. 여기서 대표값은 지정하기 나름이지만 이 문제에서는 집합 내에서 가장 작은 값이 해당 집합의 대표값이 된다고 하겠다. 그래서 매개변수로 받아온 원소와 그 원소의 parent 값이 다르다면 그 원소는 해당 집합의 대표값이 아니므로 대표값을 찾기 위해 한번 더 find를 호출하게 된다. 이를 구현한 코드는 다음과 같다.

public static int find(int a) {
    if (parent[a] == a) {
        return a;
    } else {
        return parent[a] = find(parent[a]);
    }
}

그 다음 합집합 연산을 구현한다. 합집합 연산은 매개변수로 받은 두 원소의 대표값을 비교하여 값이 서로 다르다면 이 둘을 연결해주는 역할을 한다. 두 원소의 대표값이 다르다는 것은 이 둘이 속한 집합이 다르다는 것을 의미한다. 따라서 find 값으로 대표값을 찾고 두 값이 다르다면 더 작은 대표값을 가지는 집합에 큰 대표값을 가지는 집합을 합쳐주면 된다. 여기서 더 작은 대표값으로 바꿔주는 이유는, 현재 대표값을 집합 내에서 가장 작은 값으로 정했기 때문이다. 결과적으로 대표값이 큰 집합의 값을 바꾸어주면 된다. 

설명이 애매할 수 있지만 코드를 보면 꽤 직관적이다. a 원소가 속한 집합의 대표값은 aParent, b는 bParent로 표현한다. 둘 중에서 더 큰 값의 parent 값을 더 작은 값으로 수정해준다. 이제는 a, b를 find하면 두 값이 같은 값을 가지게 되고 이를 두 원소가 연결되어 있다고 판단하면 된다. 

public static void union(int a, int b) {
    int parentA = find(a);
    int parentB = find(b);

    if (parentA < parentB) {
        parent[parentA] = parentB;
    } else if (parentA > parentB) {
        parent[parentB] = parentA;
    }
}

union과 find 연산을 구현했다면 이제 문제 조건에 맞게 코드를 작성하면 된다.

코드

import java.util.*;
import java.io.*;

public class Main {
    static int N, M;
    static int[] parent;

    public static void main(String[] args) throws Exception {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());

        N = Integer.parseInt(st.nextToken());
        M = Integer.parseInt(st.nextToken());

        parent = new int[N + 1];

        for (int i = 1; i <= N; i++) {
            parent[i] = i;
        }

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

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

    public static void union(int a, int b) {
        int parentA = find(a);
        int parentB = find(b);

        if (parentA < parentB) {
            parent[parentA] = parentB;
        } else if (parentA > parentB) {
            parent[parentB] = parentA;
        }
    }

    public static int find(int a) {
        if (parent[a] == a) {
            return a;
        } else {
            return parent[a] = find(parent[a]);
        }
    }
}