[백준/Python] 1717. 집합의 표현

 

• 같은 집합에 포함되어있는지 확인하려고 할 때 x와 y가 같은 보스를 가지고 있는지 확인하면 되기 때문에 유니온 파인드를 이용해서 풀었다
• flag를 통해 0일 때는 x,y를 유니온 1일 때는 x,y의 보스가 같은지 확인하였다.

import sys
input = sys.stdin.readline # 입력시간을 줄이기 위한 코드
sys.setrecursionlimit(100000) # 재귀를 더 깊게 들어가기 위한 코드

N,M = map(int,input().split())
arr = [-1] * (N+1)

for i in range(N+1):
    arr[i] = i # 리스트를 자기 자신의 번호로 초기화

def findboss(member):

    if arr[member] == member: # 자기 자신이 보스라면 자신을 반환
        return member
    else:
        arr[member] = findboss(arr[member]) # 보스가 있다면 보스의 보스를 찾아서 반환
        return arr[member]

def union(x, y):

    fa, fb = findboss(x), findboss(y)

    if fa == fb: # x와 y의 보스가 같다면 리턴
        return
    else:
        arr[fb] = fa # 보스가 같지 않다면 y의 보스는 x
        return


for i in range(M):
    flag,x,y = map(int,input().split())

    if flag == 0:
        union(x, y)
    else:
        if findboss(x) == findboss(y): # x와 y의 보스를 비교
            print('YES')
        else:
            print('NO')

 

findboss함수를 원래는 이런식으로 작성했었는데

아래와 같이 코드를 수정하니까 시간초과를 해결할 수 있었다.