코딩 테스트 [그래프] - 집합 표현하기

초기에 {0}, {1}, ... {n} 이 각각 n + 1 개의 집합을 이루고 있다. 여기에 합집합 연산과 두 원소가 같은 집합에 포함돼 있는지를 확인하는 연산을 수행하려고 한다. 집합을 표현하는 프로그램을 작성하시오.

(시간 제한 2초)

---

입력

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이고, 같을 수도 있다.

출력

1로 시작하는 입력에 1줄에 1개씩 YES 또는 NO로 결과를 출력한다.

예제 입력
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

예제 출력
NO
NO
YES

---

1단계 - 문제 분석하기

최대 원소의 개수 1,000,000과 질의 개수 100,000이 큰 편이므로 경로 압축이 필요한 정형적인 유니온 파인드 문제이다.

2단계 - 손으로 풀어 보기

1 처음에는 노드가 연결돼 있지 않으므로 각 노드의 대표 노드는 자기 자신이다. 각 노드의 값을 자기 인덱스값으로 초기화한다.

2 find 연산으로 특정 노드의 대표 노드를 찾고, union 연산으로 2개의 노드를 이용해 각 대표 노드를 찾아 연결한다. 그리고 질의한 값에 따라 결과를 반환한다.

3단계 - sudo코드 작성하기

N(원소 개수) M(질의 개수)
parent(대표 노드 저장 배열)

for(N만큼 반복)
	대표 노드를 자기 자신으로 초기화

for(M만큼 반복)
{
	if(0이면)
    	집합 합치기 -> union 연산
    else
    	같은 집합인지 확인하고 결괏값 출력
}

// union 연산
union(a, b)
{
	a와 b의 대표노드 찾기
    두 원소의 대표 노드끼리 연결하기
}

// find 연산
find(a)
{
	if(a가 대표노드이면)
    	리턴
    else
    	a의 대표 노드값을 find(parent[a]) 값으로 저장
}

// 두 원소가 같은 집합인지 확인
checkSame(a, b)
{
	a와 b의 대표 노드 찾기
    if(두 대표 노드가 같으면)
    	return true
    else
    	return false
}

4단계 - 코드 구현하기

import java.util.Scanner;

public class Q50 {
    public 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 question = sc.nextInt();
            int a = sc.nextInt();
            int b = sc.nextInt();
            if(question == 0)
                union(a, b);
            else{
                if(checkSame(a, b)){
                    System.out.println("YES");
                }
                else {
                    System.out.println("NO");
                }
            }
        }
    }

    public static void union(int a, int b){
        a = find(a);
        b = find(b);
        if(a != b){
            parent[b] = a;
        }
    }

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

    public static boolean checkSame(int a, int b){
        a = find(a);
        b = find(b);
        if(a == b)
            return true;
        else
            return false;
    }
}

---

• Do it! 알고리즘 코딩테스트 자바 편