[PS] BOJ 1976 / 여행 가자
문제 링크: 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