[BOJ] 1717 : 집합의 표현 [유니온파인드](Python)

문제 링크

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

1717번: 집합의 표현

소스 코드

import sys
sys.setrecursionlimit(10**6)

def main():
    n, m = map(int, input().split())
    parents = [x for x in range(n+1)] # 부모 노드 저장

    def find(x): # 부모 찾기
        if parents[x] != x: # 자기자신과 같아질 때까지
            parents[x] = find(parents[x])
        return parents[x]

    def union(a, b): # a,b 합치기
        a, b = find(a), find(b)
        if a < b: # 작은 숫자가 부모
            parents[b] = a
        else:
            parents[a] = b

    def connected(a, b): # 부모 같은지 확인
        return find(a) == find(b)

    for _ in range(m):
        x, a, b = map(int, sys.stdin.readline().strip().split())
        if x == 0:
            union(a, b)
        else:
            print("YES" if connected(a, b) else "NO")

if __name__ == '__main__':
    main()

 

해설

• 유니온 파인드

최근 코세라에서 알고리즘 강의를 들으며 유니온 파인드에 대해 복습하게 되었습니다.

로직과 구현 자체가 모두 간단한 편입니다.

아쉬웠던 것은 유니온 파인드 알고리즘 중에 효율성이 가장 좋은 weighted 방식을 적용할 수 있는 문제가 아니었다는 점입니다.

대신 굉장히 전형적인 유형의 문제였기 때문에 그 풀이를 정리하고자 합니다.

---

부모 찾기

def find(x): # 부모 찾기
    if parents[x] != x: # 자기자신과 같아질 때까지
        parents[x] = find(parents[x])
    return parents[x]

예시를 통해 이해하는 것이 가장 쉬울 것 같습니다.

0부터 7까지의 노드를 가진 그래프가 위와 같다고 합시다.

이때 0,1,2,3,4는 전부 0을 중심으로 연결되어 있고, 5와 6은 따로, 그리고 7은 혼자 구분되어 있습니다.

따라서 각 노드의 부모(루트)가 누구인지를 기록한 parents 리스트의 값은 0, 0, 0, 0, 0, 5, 5, 7이 됩니다.

 

이때의 특징을 생각해보면 최상단의 부모 노드들은 자기 자신이 곧 부모 노드입니다.

0, 5, 7의 부모 노드가 각각 0, 5, 7 자신인 것이죠.

 

따라서 어떤 노드 x의 부모 노드를 찾고 싶다면 재귀함수를 통해 x와 그 부모 노드가 동일한 조건이 되어야 합니다.

예를 들어 x가 4인 경우, 이것의 부노 노드는 1이었습니다.

따라서 4 != 1 이므로 1의 부모 노드를 추가로 탐색합니다.

1의 부모 노드는 0이고 1 != 0 이므로 0의 부모 노드를 추가로 탐색합니다.

드디어 0과 0의 부모 노드 0은 0 == 0 이므로 재귀가 종료됩니다.

 

Union

def union(a, b): # a,b 합치기
    a, b = find(a), find(b)
    if a < b: # 작은 숫자가 부모
        parents[b] = a
    else:
        parents[a] = b

합치는 것은 일반적으로 작은 노드를 부모로 삼는다는 전제를 깔고 있습니다.

예를 들어 1과 2를 합치면 2가 1에 속하게 된다는 뜻입니다.

단, 노드를 그대로 쓰는게 아니라 부모 노드가 무엇인지를 찾아 비교해야 합니다.

 

만약 두 노드의 부모 노드가 이미 동일하다면 굳이 연산할 필요도 없긴 합니다.

그러나 두 부모 노드가 다른 경우 숫자가 큰 것이 작은 것의 밑으로 들어가는 구조입니다.

제가 위에서 그려 놓은 트리 구조를 상기하시면 좋겠네요.

 

연결 여부

def connected(a, b): # 부모 같은지 확인
    return find(a) == find(b)

뭐 사실 이건 굳이 이렇게 할 필요가 없긴 합니다만.. 그냥 관습적으로 추가해봤습니다.

두 노드의 부모가 동일하면 true, 그렇지 않으면 else가 return 됩니다.

이에 따라서 print되는 내용을 정해주면 됩니다.