[Baekjoon] 집합의 표현(python)

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

문제 접근

분리집합 문제이다.

 

분리집합문제는 Union-Find를 이용하여 해결할 수 있다.

 

이 문제는 문제 입력에 따라 Union연산을 하거나 같은 집합인지 확인하면 되는 문제였다.

 

주의할 점은 find 함수 작성할 때 최적화를 하지 않으면 시간 초과가 발생할 수 있다.

---

풀이 과정

1. 합집합 연산

def union(u, v):
    u = find(u)
    v = find(v)

    if u < v:
        parent[v] = u
    else:
        parent[u] = v

 

2. 최상위 부모 검색

최상위 부모를 검색할 때 루트가 제대로 저장되지 않은 경우는 검색하면서 다시 저장해주어야 시간 초과를 해결할 수 있다.

 

그냥 재귀를 호출하여 최상위 부모를 검색하면 시간 초과가 발생할 수 있다.

def find(v):
    if v == parent[v]:
        return v

    parent[v] = find(parent[v])

    return parent[v]

 

시간 초과 코드

def find(v):
    if v == parent[v]:
        return v

    return find(parent[v])

---

정답 코드

"""
    백준 1717 집합의 표현
"""

import sys

sys.setrecursionlimit(10**5)


def union(u, v):
    u = find(u)
    v = find(v)

    if u < v:
        parent[v] = u
    else:
        parent[u] = v


def find(v):
    if v == parent[v]:
        return v

    parent[v] = find(parent[v])

    return parent[v]


n, m = list(map(int, sys.stdin.readline().split()))
parent = [i for i in range(n + 1)]

for _ in range(m):
    cmd, a, b = list(map(int, sys.stdin.readline().split()))

    if cmd == 0:
        union(a, b)
    elif cmd == 1:
        print("YES" if find(a) == find(b) else "NO")