[백준 / 그래프] 1717 : 집합의 표현 (python)

문제 설명

예제 입력 1

7 8
0 1 3
1 1 7
0 7 6
1 7 1
0 3 7
0 4 2
0 1 1
1 1 1

예제 출력 1

NO
NO
YES

 

parent 배열 초기 상태

1	2	3	4	5	6	7
1	2	3	4	5	6	7

 

union(1, 3)

1	2	3	4	5	6	7
1	2	1	4	5	6	7

 

1 1 7   # NO

 

union(7, 6)

1	2	3	4	5	6	7
1	2	1	4	5	7	7

 

1 7 1   # NO

 

union(3, 7)

1	2	3	4	5	6	7
1	2	1	4	5	7	1

 

union(4, 2)

1	2	3	4	5	6	7
1	4	1	4	5	7	1

 

union(1, 1)

1	2	3	4	5	6	7
1	4	1	4	5	7	1

1 1 1   # YES

 

풀이

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

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

def find(a):
    if a == parent[a]:
        return a
    else:
        parent[a] = find(parent[a])
        return parent[a]

def union(a, b):
    a = find(a)
    b = find(b)
    if a != b:
        parent[b] = a

for i in range(m):
    question, a, b = map(int, input().split())
    if question == 0:
        union(a, b)
    elif question == 1:
        if find(a) == find(b):
            print("YES")
        else:
            print("NO")

 

유니온 파인드 문제

유니온과 파인드를 함수로 구현하면 쉽게 풀 수 있다.