[SW Expert Academy] 3289: 서로소 집합 (Java) / Union Find

문제

SW Expert Academy

설명 (최종 코드는 아래)

서로 같은 집합에 속해있는지 확인할 때 Union Find 알고리즘을 이용할 수 있다.

int 배열로 트리 형태를 구현해 알아보는 원리이다.

핵심 함수는 3가지 이다.

1. 초기화

int[] parent = new int[size]; // 숫자는 1~ size-1까지

for(int i=0; i< size; i++)
	parent[i] = i;

맨 처음에는 자기 자신을 갖도록 한다. (size == 6일 때)

2. Union (합치기)

private static void union(int a, int b) {
		a = findSet(a);
		b = findSet(b);

		if (a != b) {
			if (a < b) {
				parent[b] = a;
			} else {
				parent[a] = b;
			}
		}
}

union(2,3) 을 실행하면 2와 3은 같은 집합에 속하게 된다. 이때, 2와 3 중 더 큰값이 작은값을 부모로 갖게 한다. (parent[3] =2) 

그림으로 나타내면 아래와 같다.

이 때의 배열 모습이다.

다음으로, union(1,2)와 union(4,5) 까지 해보자

1,2,3 이 한 집합에 묶이고 4,5 가 한 집합에 묶인다.

 

3. findSet(최상단 부모 찾기)

	private static int findSet(int x) {
		if (parent[x] == x)
			return x;
		else
			return parent[x] = findSet(parent[x]);
	}

 

위 그림과 같은 상태에서 findSet(3) 을 해주면 

1. else 로 들어가 parent[3] = findSet(parent[3]); 를 재귀 호출한다.

2. findSet(2)가 호출된다. else 로 들어가 parent[2] = findSet(parent[2]); 를 재귀 호출한다.

3. findSet(1)이 호출된다. if문 안으로 들어가 1을 반환한다. 

4. 최종적으로 배열은 아래의 모습이되고, 1이 반환된다. 

 

최종 코드

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

public class Solution {
	static BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
	static BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
	static StringTokenizer st;
	static int[] parent;
	static int T, size, M;

	public static void main(String[] args) throws Exception {
		T = Integer.parseInt(br.readLine());
		int oper, a, b;
		for (int t = 1; t <= T; t++) {
			bw.append("#" + t + " ");
			st = new StringTokenizer(br.readLine());
			size = Integer.parseInt(st.nextToken());
			M = Integer.parseInt(st.nextToken());
			parent = new int[size + 1];
			for(int i=0; i<= size; i++) {
				parent[i] = i;
			}
			for (int i = 0; i < M; i++) {
				st = new StringTokenizer(br.readLine());
				oper = Integer.parseInt(st.nextToken());
				a = Integer.parseInt(st.nextToken());
				b = Integer.parseInt(st.nextToken());
				if (oper == 0)
					union(a, b);
				if (oper == 1) {
					if (findSet(a) == findSet(b)) {
						bw.append(1 + "");
					} else {
						bw.append(0 + "");
					}
				}
			}
			bw.append("\n");
		}
		bw.flush();
		bw.close();
	}

	private static int findSet(int x) {
		if (parent[x] == x)
			return x;
		else
			return parent[x] = findSet(parent[x]);
	}

	private static void union(int a, int b) {
		a = findSet(a);
		b = findSet(b);

		if (a != b) {
			if (a < b) {
				parent[b] = a;
			} else {
				parent[a] = b;
			}
		}
	}
}