[PS] BOJ 13911 / 집 구하기

[PS] BOJ 13911 / 집 구하기
Thumbnail: Photo by Vladimir Wang (Unsplash)
문제 링크: https://www.acmicpc.net/problem/13911

다익스트라 알고리즘을 활용해 최단 거리를 구하는 문제입니다. 다만, 약간의 테크닉을 요구합니다.

풀이

다익스트라 알고리즘

​다익스트라는 주어진 그래프에서 두 정점 사이의 최단 거리를 구하는 알고리즘입니다. 임의의 한 정점에서 시작해, 가장 거리가 가까운 정점을 찾아가며 전체 정점의 최단 거리를 갱신하는 방식으로 동작합니다. 자세한 설명은 이전 글을 참고해 주세요.

풀이

문제에서, 총 $V$개의 정점과 $E$개의 간선이 주어집니다. 문제에서 간선을 도로라고 언급하므로, 양방향 그래프로 간주하고 그래프를 구성했습니다.

전체 정점 중 $M$($1 \leq M \leq V - 2)개는 맥도날드, $S$($1 \leq S \leq V - 2$) 개는 스타벅스입니다. 나머지 정점은 모두 집입니다.

우리가 찾고자 하는 집은 다음 3가지 조건을 모두 만족해야 합니다.

  • 어떤 집이 '맥세권'이려면, 가장 가까운 맥도날드와 집 사이의 거리가 $x$ 이하여야 합니다.
  • 어떤 집이 '스세권'이려면, 가장 가까운 스타벅스와 집 사이의 거리가 $y$ 이하여야 합니다.
  • 맥세권과 스세권을 모두 만족하는 집 중, 집에서부터 맥도날드까지의 거리와 스타벅스까지의 거리의 합이 최소인 집

다익스트라는 1개의 정점에서 나머지 모든 정점으로의 최단 거리를 계산합니다. 하지만 이걸 $M + S$개의 정점에서 각각 수행하게 되면 문제의 제한 시간보다 더 많은 시간이 필요합니다. 다익스트라의 실행 횟수를 줄일 수 있는 방법이 필요합니다.

가상의 정점을 활용해 다익스트라 횟수 줄이기

$M$개의 맥도날드와 $S$개의 스타벅스에서 다익스트라를 각각 진행하기보다, 가상의 정점에 $M$개의 맥도날드와 $S$개의 스타벅스를 연결한 뒤, 가상의 정점에서부터 다익스트라를 진행한다면 다익스트라가 진행되는 횟수를 줄일 수 있습니다.

만약 $M+S$번 다익스트라를 진행한다면, 다익스트라의 반복 횟수도 문제일 뿐 더러 각각의 다익스트라로부터 구한 최단 거리를 비교해 최종 거리를 계산하는 과정도 필요합니다.

가상의 정점을 2개 만들고, 한 개의 정점에는 모든 맥도날드를, 다른 한 개에는 모든 스타벅스를 연결한 뒤 2개의 정점에서 다익스트라를 각각 진행합니다.

가상의 정점 1개에 모두 연결하지 않는 이유는, 맥도날드와의 최단 거리와 스타벅스와의 최단 거리가 모두 필요하기 때문입니다.

from heapq import heappush, heappop
INF = 100_000_001 # 도달할 수 없는 정점을 의미하는 값.

# 입력 처리
V, E = map(int, input().split())
graph = tuple([] for _ in range(V + 3)) # 1 ~ V의 정점과 2개의 가상의 정점을 포함한다
for _ in range(E):
    u, v, w = map(int, input().split())
    graph[u].append((v, w))
    graph[v].append((u, w))

M, x = map(int, input().split())
mcdonalds = tuple(map(int, input().split()))
for m in mcdonalds:
    graph[V + 1].append((m, 0)) # 가상의 정점 V + 1에 모든 맥도날드 정점을 연결

S, y = map(int, input().split())
starbucks = tuple(map(int, input().split()))
for s in starbucks:
    graph[V + 2].append((s, 0)) # 가상의 정점 V + 2에 모든 스타벅스 정점을 연결

# 우선순위 큐 기반의 다익스트라 구현
queue = []
def dijkstra(src):
    distance = [INF] * (V + 3)
    distance[src] = 0
    queue.clear()
    heappush(queue, (0, src))
    while queue:
        c, u = heappop(queue)
        if c > distance[u]:
            continue

        for v, w in graph[u]:
            if c + w < distance[v]:
                distance[v] = c + w
                heappush(queue, (distance[v], v))
            
    return distance

# 맥도날드와 스타벅스를 같은 정점에 연결할 경우, 한 집에서 맥도날드와 스타벅스의 구분 없이 최단 거리가 계산되므로 각각은 다른 정점에 연결해야 한다.
mcdonalds_dist = dijkstra(V + 1)
starbucks_dist = dijkstra(V + 2)

가상의 정점을 활용해 다익스트라 횟수 줄이기

2번의 다익스트라 이후에는 문제 조건에 맞게 맥세권과 스세권에 해당하는 집 중 두 최단 거리의 합이 최소가 되는 경우를 찾으면 됩니다.

전체 코드

from heapq import heappush, heappop
INF = 100_000_001
input = open(0).readline
V, E = map(int, input().split())
graph = tuple([] for _ in range(V + 3))
for _ in range(E):
    u, v, w = map(int, input().split())
    graph[u].append((v, w))
    graph[v].append((u, w))

M, x = map(int, input().split())
mcdonalds = tuple(map(int, input().split()))
for m in mcdonalds:
    graph[V + 1].append((m, 0)) # 더미 노드인 V+1번에 맥도날드 정점을 모두 연결해 한 번의 다익스트라로 최단거리 계산 처리

S, y = map(int, input().split())
starbucks = tuple(map(int, input().split()))
for s in starbucks:
    graph[V + 2].append((s, 0)) # 더미 노드인 V+2번에 스타벅스 정점을 모두 연결해 한 번의 다익스트라로 최단거리 계산 처리

queue = []
def dijkstra(src):
    distance = [INF] * (V + 3)
    distance[src] = 0
    queue.clear()
    heappush(queue, (0, src))
    while queue:
        c, u = heappop(queue)
        if c > distance[u]:
            continue

        for v, w in graph[u]:
            if c + w < distance[v]:
                distance[v] = c + w
                heappush(queue, (distance[v], v))
            
    return distance

# 맥도날드와 스타벅스를 같은 더미 노드에 연결할 시, 한 집에서 맥도날드와 스타벅스의 구분 없이 최단 거리가 계산되므로 각각은 다른 더미 노드로 연결해야 한다.
mcdonalds_dist = dijkstra(V + 1) # 더미 노드 통해서 모든 맥도날드로의 최단 거리 계산하기
starbucks_dist = dijkstra(V + 2) # 더미 노드 통해서 모든 스타벅스로의 최단 거리 계산하기

ans_sum = INF
for i in range(1, V + 1):
    if mcdonalds_dist[i] > x or starbucks_dist[i] > y or mcdonalds_dist[i] == 0 or starbucks_dist[i] == 0: # 맥세권 및 스세권인 집만 해당한다.
        continue
    this_sum = mcdonalds_dist[i] + starbucks_dist[i]
    if this_sum < ans_sum:
        ans_sum = this_sum

print(ans_sum if ans_sum < INF else -1)

solution.py