집합의 표현 1717번 python

Union-Find를 이용한 풀이:

import sys
sys.setrecursionlimit(10**5)

def union(x,y):
    
    x=union_find(x)
    y=union_find(y)
    
    if x<y:
        parent[y]=x
    else:
        parent[x]=y
        
def union_find(x):
    
    if parent[x]==x:
        return x
    parent[x]=union_find(parent[x])
    return parent[x]
    

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

parent=[0]*(n+1)

for i in range(1,n+1):
    parent[i]=i

for i in range(m):
    x,y,z=map(int,input().split())
    
    if x==0:
        union(y,z)
        print(parent)
    else:
        
        if union_find(y)==union_find(z):
            print('YES')
        else:
            print('NO')

기존에 했던방식

 

def find_union(n):
    
    if parent[n]!=n:
        return find_union(parent[n])
    return n

 

더 좋은 코드

def union_find(x):
    
    if parent[x]==x:
        return x
    parent[x]=union_find(parent[x])
    return parent[x]

 

해당하는 방식으로 해야 parent값이 계속 업데이트가 된다. 주의하자