[Python] 백준 / gold / 1976 : 여행 가자

문제 링크 : https://www.acmicpc.net/problem/1976

Union-Find 문제. 서로소 집합, 분리 집합 문제라고도 한다. 서로 같은 집합 내에 속하는지, 다른 집합인지 구분하는 문제다. 트리 자료구조에는 루트 노드가 존재한다. Union-Find 문제에선, 같은 집합내의 원소들을 모아 트리를 짓고, 루트 노드가 같은지 여부를 통해 같은 집합 내에 속하는지를 판단한다.

이 유형을 처음 접했을때는 "루트 노드 테이블을 갱신해서 푼다" 라고만 기억해뒀는데, 오늘 정확한 개념을 잡았다. 말 그대로 "Union 과 Find 메서드를 명확히 구현해주는 것"이 좋다.

find 메서드는 정말 루트 노드를 찾아가는 메서드로만 구현해야 하고, Union 메서드는 집합을 합쳐주는 합집합 메서드로 구현해야 한다.

2가지 기억할 것
1. Union-Find 를 명확히 구현한다.
2. 루트끼리 합쳐서 합집합 연산을 구현한다.

파이썬 코드

import sys
N = int(sys.stdin.readline())
M = int(sys.stdin.readline())

graph = []
for _ in range(N):
    graph.append(list(map(int, sys.stdin.readline().split())))

plan = list(map(int, sys.stdin.readline().split()))
parent = [x for x in range(N)]


# find
def findParent(x):
    if parent[x] == x:
        return x
    else:
        parent[x] = findParent(parent[x])
        return parent[x]


# union
for i in range(N):
    for j in range(N):
        if graph[i][j] == 1:
            pi = findParent(i)
            pj = findParent(j)
            if pi < pj:
                parent[pj] = pi
            else:
                parent[pi] = pj

#print(parent)
root = parent[plan[0]-1]
for p in plan:
    if root != parent[p-1]:
        print("NO")
        sys.exit(0)

print("YES")