백준 1717 집합의 표현 python

CS 스터디에서 union-find를 학습한 김에 풀어보았다.

해당 알고리즘의 원리를 다시 되새기면서 구현 연습을 계속 해봐야겠다.

N, M = map(int, input().split())
p = [0] * (N+1) # 부모 리스트
r = [0] * (N+1) # rank 리스트

def makeSet():
    for i in range(0,N+1):
        p[i] = i # 맨 처음엔 union되지 않았으므로. 자기 자신을 부모로 하는 집합을 생성합니다.
        r[i] = 1

def union(x,y):
    x = find(x)
    y = find(y) # union을 하려면 각 원소의 부모를 알아야 합니다.

    if x == y: # 두 원소의 부모가 같다는 뜻은, 같은 집합이란 뜻입니다.
        return False # 그러므로 union 불가입니다.

    if r[x] < r[y]: # 만약 y원소의 가중치가가 더 높다면
        r[y] += r[x] # y원소가 부모일 확률이 높게 만들어야 하므로 y 가중치 +
        p[x] = y # y원소에 x를 붙일거임
    else:
        r[x] += r[y]
        p[y] = x
    return True

def find(x):
    if x == p[x]: # 제일 윗 부모를 찾았다면 (자기 자신을 가리키고 있다면)
        return x # 그 값 return
    else:
        #주의! find의 인자로 x를 넘기게 되면, 계속 같은 과정을 반복합니다. 잘 외우세요.
        p[x] = find(p[x]) # 부모를 찾지 못했다면,
        return p[x] # 부모를 찾아 return

makeSet()

union(0, 0)

for _ in range(M):
    cal, a, b = map(int, input().split())
    if cal == 0:
        union(a, b)
    else:
        if find(a) == find(b):
            print('YES')
        else:
            print('NO')