Baekjoon Online Judge

algorithm practice

- 단계별 문제풀기

29. 유니온 파인드

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

Java / Python

2. 여행 가자

1976번

BFS, DFS 뿐만 아니라 유니온 파인드로도 두 정점이 연결되어 있는지를 확인할 수 있습니다.

이번 문제는 도시들의 개수와 도시들 간의 연결 여부가 주어져 있고, 동혁이의 여행 계획에 속한 도시들이 순서대로 주어졌을 때 가능한지 여부를 판별하는 프로그램을 작성하는 문제이다. 같은 도시를 여러 번 방문하는 것도 가능하다.

parent를 자기 자신으로 초기화하고, 부모를 찾는 함수 find와 합집합 연산을 해, 같은 부모를 가지도록 하는 union함수를 이용한다.

• 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;

		int N = Integer.parseInt(br.readLine());
		int M = Integer.parseInt(br.readLine());

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

		for (int i = 1; i <= N; i++) {
			st = new StringTokenizer(br.readLine());
			for (int j = 1; j <= N; j++) {
				int map = Integer.parseInt(st.nextToken());

				// 연결된 부분은 합집합 연산함.
				if (map == 1) {
					union(i, j);
				}
			}
		}
		st = new StringTokenizer(br.readLine());
		int start = find(Integer.parseInt(st.nextToken()));
		for (int i = 1; i < M; i++) {
			int idx = Integer.parseInt(st.nextToken());

			if (start != find(idx)) {
				bw.write("NO\n");
				bw.flush();
				bw.close();
				br.close();
				return;
			}
		}

		bw.write("YES\n");
		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;
			}
		}
	}
}

• Python

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

N = int(input()) # 노드 수 
M = int(input()) # 정점 수 
parent = [i for i in range(N+1)]

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: # 동일한 집합일 경우
        return

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

for y in range(1, N+1): 
    maps = list(map(int, input().split())) 
    for x in range(1, len(maps)+1): # y 도시 연결 정보 추출
        if maps[x-1] == 1: # x 도시와 연결되어 있으면
            union(y, x) 
check = True
tour = list(map(int, input().split())) # 여행 계획 정보
result = set([find(i) for i in tour]) 
if len(result) != 1: # set 2개 이상이면 두 개의 집합 존재
    print('NO')
else:
    print('YES') # set 1이면 한 개의 집합 존재