[백준] 1717번 집합의 표현 / 파이썬 Python (union find)

http://https://www.acmicpc.net/problem/1717

 

 

백준 1717번 집합의 표현은 유니온 파인드(union find) 알고리즘을 확실히 알게 해주는 문제였다.

 

문제

 

 

설명

 

문제 1717번 예제로 유니온 파인드에 대해 설명하자면,

 

1)  대표 노드를 저장하는 1차원 리스트가 필요하며 처음에는 연결되어 있지 않으므로

자신의 인덱스 값으로 초기화한다.

 

 

2) 2개의 노드 각각의 대표 노드를 찾아 union 연산으로 연결한다.

예를 들어 두 번째 줄 입력을 받아서 0일 때 union(1,3) 을 해주면 arr[3] = 1로 값이 변경된다.

이때, 작은 값을 대표 노드로 본다.

 

유니온 파인드 알고리즘에서 유의할 것은 입력 6번째 줄 union(3,7) 인 경우다.

5번째 줄까지 모두 진행되었다고 가정해보면 union(7,6)이 실행되었으므로 arr[7] = 6 이다.

 

3, 7번의 각 대표 노드를 찾은 후, 그 대표 노드끼리 연결해야 한다.

3번의 대표 노드는 1, 7번의 대표 노드는 6 이므로 arr[7] = 1 이 아닌 arr[6] = 1 로 변경해준다.

 

 

 

3) find 연산은 자신의 대표 노드를 찾아주는 역할이다.

index와 value가 같을 때까지 찾아야 하므로 재귀 함수로 구현한다.

 

 

코드

import sys
sys.setrecursionlimit(1000000)
input = sys.stdin.readline

n, m = map(int, input().split())
arr = [i for i in range(n+1)]

def Find(f):				# 대표 노드 찾기
    if arr[f] != f:
        arr[f] = Find(arr[f])	# 재귀로 구현
    return arr[f]

def Union(x, y):			# 대표 노드끼리 합치는 union 연산
    idx = Find(x)
    idx2 = Find(y)
    if idx < idx2:			# 작은 값을 대표 노드로
        arr[idx2] = idx
    else:
        arr[idx] = idx2

for _ in range(m):
    flag, a, b = map(int, input().split())
    if flag == 0:
        Union(a, b)
    else:
        if Find(a) != Find(b):
            print("NO")
        else:
            print("YES")