골드5

문제 1717

해당 문제 링크

정답

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

def find(x):
    if parent[x] != x:
        parent[x]= find(parent[x])
    return parent[x]

def union(a, b):
    a=find(a)
    b=find(b)
    #값이 더 작은 것이 부모 노드로 취급함.
    if a<b:
        parent[b]=a
    else:
        parent[a]=b

v, e = map(int, input().split())
# parent=[0]* (v+1)
# for i in range(1, v+1):
#     parent[i]=i

parent = [i for i in range(v+1)]


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

풀이

시간 초과가 많이 난 문제였다.
시간 초과가 나지 않도록 코드를 작성할 때 주의할 부분이 있다!!

1. 재귀를 사용해서 풀어야 하는 문제라면 이렇게 쓰는 것은 필수이다.

파이썬의 기본 재귀 깊이 제한은 1000으로 매우 얕은 편이다.
따라서 재귀로 문제를 풀 경우 드물지 않게 이 제한에 걸리게 된다. 이 함정에 걸린 사람은 원인 미상의 런타임 에러를 붙잡고 있게된다.

import sys
sys.setrecursionlimit(10**6)

2. 반복문으로 여러줄 입력받는 상황에서는 반드시 sys.stdin.readline()을 사용해야 시간초과가 발생하지 않습니다.

input = sys.stdin.readline

3. find 함수 부분을 '경로 압축' (Path Compression) 을 통해 다음과 같이 리팩토링 할 수 있다.

시간 복잡도가 감소하므로 이 코드를 사용하는 것이 좋다.

def find(x):
    if parent[x] != x:
        parent[x]= find(parent[x])
    return parent[x]

4. for문을 한줄로 표현

위의 코드보다 아래와 같이 바꾸는 것이 더 빠르다.

parent=[0]* (v+1)
for i in range(1, v+1):
    parent[i]=i

parent = [i for i in range(v+1)]