[PS] BOJ 1976 / 여행 가자

[PS] BOJ 1976 / 여행 가자
Thumbnail: Generated using thumbnail maker
문제 링크: https://www.acmicpc.net/problem/1976

분리 집합 (Disjoint Set) 기반의 유니온 파인드 알고리즘을 활용하는 문제입니다.

풀이

​유니온 파인드(Union Find)

문제에서는 인접 행렬로 임의의 두 도시 사이의 연결 관계를 제공합니다. 가중치는 따로 없고, 두 도시 $i$와 $j$가 연결되어 있다면, 반대로도 연결되어 있습니다. (무방향)

이런 조건에서, 임의의 두 정점이 연결되었는지를 판단할 때 유니온 파인드를 사용할 수 있습니다.

분리 집합을 활용한 유니온 파인드 알고리즘의 자세한 설명은 이전 글을 참고해주세요.

N = int(input())
M = int(input())

# Disjoint Set (Union-Find)
parent = [i for i in range(N + 1)]
def find(x):
    path = []
    while x != parent[x]:
        path.append(x)
        x = parent[x]
    for p in path: # parent배열을 평탄화해, 되도록 탐색의 깊이가 얕게 유지되게 한다.
        parent[p] = x
    return x

def union(a, b):
    a = find(a)
    b = find(b)
    if a != b:
        parent[b] = a
        return True
    return False

분리 집합 기반의 유니온 파인드 구현

인접 행렬로 주어진 두 도시 간의 연결 관계를 저장할 때, 유니온 파인드를 진행해 연결 관계를 파악합니다.

이제와서 생각해보니, 굳이 graph배열을 만들어 저장할 필요는 없었네요..
graph = [[0 for _ in range(N + 1)] for _ in range(N + 1)]

for i in range(1, N + 1):
    for j, connected in enumerate(map(int, input().split()), start=1):
        graph[i][j] = connected
        if connected == 1:
            union(i, j)

입력 받기

답 출력하기

이후, M개의 도시를 거치는 여행 경로에 대해, 인접한 두 도시가 서로 연결되어있지 않다면 즉시 NO를 출력하고, 모두 연결되어있다면 YES를 출력해주면 됩니다.

plan = list(map(int, input().split()))
for v in range(1, M):
    if find(plan[v - 1]) != find(plan[v]):
        print("NO")
        break
else:
    print("YES")

답 출력하기

전체 코드

input = open(0).readline
N = int(input())
M = int(input())
graph = [[0 for _ in range(N + 1)] for _ in range(N + 1)]

# Disjoint Set (Union-Find)
parent = [i for i in range(N + 1)]
def find(x):
    path = []
    while x != parent[x]:
        path.append(x)
        x = parent[x]
    for p in path:
        parent[p] = x
    return x

def union(a, b):
    a = find(a)
    b = find(b)
    if a != b:
        parent[b] = a
        return True
    return False

for i in range(1, N + 1):
    for j, connected in enumerate(map(int, input().split()), start=1):
        graph[i][j] = connected
        if connected == 1:
            union(i, j)

plan = list(map(int, input().split()))
for v in range(1, M):
    if find(plan[v - 1]) != find(plan[v]):
        print("NO")
        break
else:
    print("YES")

solution.py