BOJ-1717 - Python

n개의 노드가 있을 때 각 노드들이 연결되어 있는 지 여부를 출력 및 합연산을 하는 유니온 파인드 문제. 기존 유니온 파인드로 풀 수 있는 문제들을 플로이드 와샬, DFSBFS 등의 방법으로 풀었으나 유니온 파인드 방법도 알아두면 좋을 것 같아 공부해보았다. 유니온 파인드는 주어진 노드 또는 집합을 합하는 Union과 노드의 루트 노드가 무엇인 지 반환하는 Find로 나뉜다. 자세한 사항은 주석을 참고. 해당 문제는 간단한 유니온 파인드의 구현으로 풀었다.

import sys
input = sys.stdin.readline

def find(node):
    # 현재 노드가 루트 노드일 시
    if node == parent[node]:
        # 루트 노드 값을 반환
        return node

    # 루트 노드가 아닐 시 재귀적으로 탐색
    p = find(parent[node])
    # 루트 노드를 현재 node 값에 저장
    parent[node] = p
    return p # 해당 루트 노드 값을 반환

    # 루트 노드가 아닐 시 탐색한 후 할당, 반환하는 과정을 경로 압축 최적화라고 한다
    # return find(parent[u])와 같이 작성시 경로 압축 최적화가 안되었다고 할 수 있는데
    # 이는 단순히 배열에 할당하지 않기 때문인가 ?
    # 아무튼 사용 이유는 한 쪽으로 치우쳐진 트리일 때 걸리는 시간이라고 한다.
    # 문제를 더 풀어가며 이해해야할 부분 같다.


def union(a, b):
    # a, b의 루트 노드를 찾는다.
    parent_a = find(a)
    parent_b = find(b)

    # 루트노드가 갖지 않을 시 b의 루트 노드에 a 루트 노드 값을 저장
    # a의 루트 노드에 b 루트 노드 값을 저장해도 된다.
    if parent_a != parent_b:
        parent[parent_b] = parent_a


n, m = map(int, input().split())
# 각 노드의 루트 노드를 저장할 배열
# 처음에는 자기 자신을 루트 노드로 생성
parent = [i for i in range(n+1)]

for _ in range(m):
    c, a, b = map(int, input().split())
    if c == 0:
        union(a, b)
    else:
        parent_a = find(a)
        parent_b = find(b)
        print("YES" if parent_a == parent_b else "NO")

Written by@hyesungoh
Learning every moment

InstagramGitHubLinkedIn