클라우드로 크고 싶은 블로그 입니다

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

1717번: 집합의 표현

 

파이썬으로 아래와 같이 작성했다.

 

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

n, m = map(int, input().split())

# 노드의 부모 노드를 저장하는 배열
# 최초 각 노드의 부모는 자기 자신
parent = [i for i in range(n+1)]

# 각 노드의 랭크를 저장하는 배열
rank = [0] * (n+1)

def find(x):
    # 경로 압축: 루트 노드를 찾을 때까지 부모 노드를 따라 이동하며 부모 노드를 업데이트
    if parent[x] != x:
        parent[x] = find(parent[x])
    return parent[x]

def union(x, y):
    # 루트 노드 찾기
    x_root = find(x)
    y_root = find(y)

    # 이미 같은 집합에 속한 경우
    if x_root == y_root:
        return

    # 랭크 기반 합집: 랭크가 낮은 노드의 루트 노드를 랭크가 높은 노드의 루트 노드의 자식으로 만듭니다.
    if rank[x_root] < rank[y_root]:
        parent[x_root] = y_root
    else:
        parent[y_root] = x_root
        # 랭크가 같으면 합쳐진 집합의 랭크를 1 증가
        if rank[x_root] == rank[y_root]:
            rank[x_root] += 1

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


'''
예제 입력
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

예제 출력
NO
NO
YES
'''