Baekjoon Online Judge

algorithm practice

- 단계별 문제풀기

29. 유니온 파인드

유니온 파인드(또는 disjoint set, 상호 배타적 집합, ...) 자료구조를 배워 봅시다.

Java / Python

1. 집합의 표현

1717번

유니온 파인드(disjoint set)에 대해 알아보는 문제

이번 문제는 집합을 표현하는 프로그램을 작성하는 문제이다.

parent를 자기 자신으로 초기화하고, 부모를 찾는 함수 find와 합집합 연산을 해, 같은 부모를 가지도록 하는 union함수를 이용한다. 0인 경우, union으로 합집합 연산을 하고 1인 경우, find를 통해 부모를 찾으며 같으면 같은 집합, 다르면 다른 집합으로 YES, NO를 맞게 출력한다.

• Java

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

public class Main {
	static int[] parent;

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

		int N = Integer.parseInt(st.nextToken());
		int 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 cmd = Integer.parseInt(st.nextToken());
			int a = Integer.parseInt(st.nextToken());
			int b = Integer.parseInt(st.nextToken());

			if (cmd == 0) {
				union(a, b);
			} else if (cmd == 1) {
				sb.append((isSameParent(a, b) ? "YES" : "NO") + "\n");
			}
		}

		bw.write(sb.toString());
		bw.flush();
		bw.close();
		br.close();
	}

	// x의 부모 찾기
	public static int find(int x) {
		if (x == parent[x])
			return x;

		return parent[x] = find(parent[x]);
	}

	// y 부모를 x 부모로 치환하기 (x > y 일 경우 반대)
	public static void union(int x, int y) {
		x = find(x);
		y = find(y);

		if (x != y) {
			if (x < y) {
				parent[y] = x;
			} else {
				parent[x] = y;
			}
		}
	}

	// x, y 부모가 같은지 확인
	public static boolean isSameParent(int x, int y) {
		x = find(x);
		y = find(y);

		if (x == y)
			return true;

		return false;
	}
}

• Python

import sys
sys.setrecursionlimit(10 ** 6)
input = sys.stdin.readline

def find(x):
    if x == parent[x]:
        return x
    
    parent[x] = find(parent[x])
    return parent[x]

def union(x, y):
    x = find(x)
    y = find(y)

    if x < y:
        parent[y] = x
    else :
        parent[x] = y


N, M = map(int, input().split())
parent = [i for i in range(N+1)]

for _ in range(M):
    cmd, x, y = map(int, input().split())

    if not cmd:
        union(x, y)

    if cmd:
        if find(x) == find(y):
            print('YES')
        else:
            print('NO')