[Java] Baekjoon1717: 집합의 표현

집합의 표현

2 초

128 MB

54678

17296

10487

28.478%

문제

초기에 {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이며 같을 수도 있다.

출력

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

 

풀이

유니온 파인드를 이용한 풀이를 진행합니다.

각각의 집합의 대표를 자기 자신으로 설정해줍니다.

 n = Integer.parseInt(st.nextToken());
 int m = Integer.parseInt(st.nextToken());

 parents = new int[n+1];

 for(int i=0; i<n+1; i++) {
     parents[i] = i; //자기 자신이 대표
 }

 

해당 집합의 대표 찾기 메소드 생성

유니온 파인드의 핵심 메소드입니다.

1차원 배열에는 해당 원소의 부모가 들어있는데 부모를 연속적으로 호출해서 자기 자신이 부모가 되는 값까지 거슬러 오릅니다.(집합 대표)

static int find(int a) { //대표찾기
    if(parents[a]==a) return a; //자기 자신이 부모라면 = 곧 이 집합의 대표

    parents[a] = find(parents[a]); // 압축: 현재 원소 a를 대표 가지에 붙여줌
    return parents[a]; //대표를 찾아 반환
}

또한 부모를 찾아 오르면서 좌표압축을 해줄 수 있습니다. (현재 가지를 떼어내서 부모가지에 붙여줌)

 

합집합 메소드

현재 원소가 포함된 집합의 대표를 찾은 다음 한 집합 대표의 부모를 다른 집합 대표로 바꿔주면 서로 합쳐집니다.

이 때 find 메소드를 기반으로 대표 번호를 찾고 

static void union(int a, int b) { //합치기
    int aRoot = find(a);
    int bRoot = find(b);
    if(aRoot == bRoot) return; //두 집합이 같은 경우 리턴

    parents[bRoot] = aRoot; // 다르면 한쪽에 합치기
}

 

서로 포함관계 확인하기

위 합집합 메소드만 안다면 거의 동일한 구성으로 확인할 수 있습니다.

static void isUnited(int a, int b) { //같으면 true

    int aRoot = find(a);
    int bRoot = find(b);

    if(aRoot == bRoot) {//두 집합이 같은 경우
        sb.append("YES"+"\n"); //스트링 빌더로 저장했습니다
        return;
    }

    sb.append("NO"+"\n");
}

 

전체코드

import java.io.*;
import java.util.*;
public class Main {
	
	static int n,parents[];
	static StringBuilder sb = new StringBuilder();
	
	static int find(int a) { //대표찾기
		if(parents[a]==a) return a;
		
		parents[a] = find(parents[a]); // 압축
		return parents[a];
	}
	
	static void isUnited(int a, int b) { //같으면 true
		
		int aRoot = find(a);
		int bRoot = find(b);
		
		if(aRoot == bRoot) {//두 집합이 같은 경우
			sb.append("YES"+"\n");
			return;
		}
		
		sb.append("NO"+"\n");
	}
	
	static void union(int a, int b) { //합치기
		int aRoot = find(a);
		int bRoot = find(b);
		if(aRoot == bRoot) return; //두 집합이 같은 경우 리턴
		
		parents[bRoot] = aRoot; // 다르면 한쪽에 합치기
	}
	
public static void main(String[] args) throws IOException{
	BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
	StringTokenizer st;
	
	 st = new StringTokenizer(br.readLine());
	 n = Integer.parseInt(st.nextToken());
	 int m = Integer.parseInt(st.nextToken());
	 
	 parents = new int[n+1];
	 
	 for(int i=0; i<n+1; i++) {
		 parents[i] = i; //자기 자신이 대표
	 }
		 
	 while(m-->0) { //m개의 입력값 받아주기
		 st = new StringTokenizer(br.readLine());
		 int command = Integer.parseInt(st.nextToken());
		 int a = Integer.parseInt(st.nextToken());
		 int b = Integer.parseInt(st.nextToken());
		 
		 switch(command) { //0은 합집합, 1은 합집합 확인 연산
		 	case 0: union(a,b); break;
		 	case 1: isUnited(a, b); break;
		 }
	 }
		 
	 System.out.println(sb.toString());
}
}