[Python/파이썬] 백준 1277번 발전소 설치

1277번: 발전소 설치

 

문제

엄청난 벼락을 맞아 많은 전선들이 끊어져 현재 전력 공급이 중단된 상태이다. 가장 심각한 문제는 1번 발전소에서 N번 발전소로 가는 중간의 전선이 끊어진 것이기에 일단 이 두 발전소를 다시 연결하는게 현재 해결해야할 첫 번째 과제이다.

발전소는 1번부터 N번까지 번호로 매겨져 2차원 격자 좌표 위에 있다. 그리고 몇몇 전선은 보존된 채 몇몇 발전소를 잇고 있다. 문제는 현재 전선과 발전소의 위치가 주어졌을 때 최소의 전선 길이를 추가로 사용하여 1번 발전소와 N번 발전소를 연결짓는 것이다. 물론 연결 짓는 중간에 다른 발전소를 거쳐갈 수 있다. 단, 안정성 문제로 어떠한 두 발전소 사이의 전선의 길이가 M을 초과할 수는 없다. 아래에 이에 대한 예를 그려놓았다.

         연결 전                          연결 후

3  . . . 7 9 . . . . .          3  . . . 7 9 . . . . .
                                          /
2  . . 5 6 . . . . . .          2  . . 5 6 . . . . . .
                                        /
1  2-3-4 . 8 . . . . .          1  2-3-4 . 8 . . . . .
   |                               |
0  1 . . . . . . . . .          0  1 . . . . . . . . .

   0 1 2 3 4 5 6 7 8 9             0 1 2 3 4 5 6 7 8 9

입력

첫 줄에는 발전소의 수 N(1 ≤ N ≤ 1,000)과 현재 남아있는 전선의 수 W(1≤ W ≤ 10,000)가 주어진다. 두 번째 줄에는 제한 길이 M(0.0 < M < 200,000.0)가 주어진다. 다음 N개의 줄에는 1번 발전소부터 N번 발전소까지 각각의 발전소의 X좌표와 Y좌표(-100,000 ≤ xi,yi  ≤ 100,000)가 차례대로 주어진다. 다음 W개의 줄에 대해 각 줄에는 두 개의 정수가 입력되어지는데 이는 현재 남아있는 전선이 잇고 있는 두 발전소를 의미한다.

출력

첫 줄에 1번 발전소와 N번 발전소를 잇는데 필요한 추가 전선 길이의 최솟값을 1000배하여 출력한다. (단, 1000배하고 난 후 나머지 소수는 반올림이 아닌 버림을 한다)

코드

import sys
input = sys.stdin.readline
from heapq import heappop, heappush
from math import sqrt, floor

n, w = map(int, input().split())
m = float(input())  # 제한 길이

plants = [[]]
for _ in range(n):  # 1~N번 발전소의 위치
    plants.append(list(map(int, input().split())))
    
def get_distance(a, b):
    x = plants[a][0]-plants[b][0]
    y = plants[a][1]-plants[b][1]
    return sqrt(x**2 + y**2)

wire = [[] for _ in range(n+1)]
for _ in range(w):
    a, b = map(int, input().split())
    # 남아있는 전선들은 추가 길이에 계산되지 않도록 0으로 설정
    wire[a].append((b, 0))
    wire[b].append((a, 0))

# 발전소들 간의 거리 계산
for i in range(1, n+1):
    for j in range(i+1, n+1):
            dist = get_distance(i, j)
            if dist <= m:   # 제한 길이 넘는지 확인
                wire[i].append((j, dist))
                wire[j].append((i, dist))

# 다익스트라
q = []
heappush(q, (0, 1)) # 1번 발전소(출발점) heapq에 넣어주기
dp = [int(1e9)] * (n+1)
dp[1] = 0
while q:
    dist, now = heappop(q)
    if dp[now] < dist:
        continue
    for nx, nd in wire[now]:
        next_dist = dist + nd
        if dp[nx] > next_dist:
            dp[nx] = next_dist
            heappush(q, (next_dist, nx))

print(floor(dp[n]*1000))

• 1번 발전소부터 N번 발전소까지의 거리를 구해야하므로 한 정점에서 다른 모든 정점까지의 최단거리를 구하는 다익스트라(Dijkstra) 알고리즘을 이용하여 풀이.
• 제한 길이를 넘지 않아서 서로 연결 가능한 발전소들을 wire 배열에 (갈 수 있는 발전소, 필요 전선 길이) 형태의 튜플들 리스트로 저장. 즉 i -> j 발전소가 서로 연결 가능하다면 wire[i]에 (j, 전선 길이), wire[j]에 (i, 전선 길이) 형태로 저장.
  이때 남아있는 전선들은 추가 전선 길이를 계산할 때 포함되지 않도록 전선 길이를 0으로 설정하여 넣어준다.
• wire 배열의 값들을 다익스트라 알고리즘을 이용하여 1번 발전소부터 다른 발전소로 가는 최단 경로를 모두 구하여 dp배열에 넣어준다.
• 문제에서 요구하는 값인 N번 발전소까지 잇는데 필요한 최소 전선의 길이는 dp[n]의 값이다.