[BOJ / Python] 13911번: 집 구하기

https://www.acmicpc.net/problem/13911

13911번: 집 구하기

 

문제

 

풀이

처음에는 맥도날드와 스타벅스가 없는 정점들을 순차적으로 출발점으로 두고 반복문을 사용하여 맥도날드와 스타벅스의 최단거리의 합을 각각 구했다. 하지만 이 풀이는 시간 초과라는 결과를 가져와 다른 방식으로 접근해야겠다고 생각했다. 먼저 맥도날드에 대해 모든 맥도날드 위치를 출발점으로 두고 최소힙을 사용하여 각 정점들에 대해 거리를 구했다. 그리고 스타벅스에 대해서도 모든 스타벅스의 위치를 출발점으로 두고 최소힙으로 각 정점들의 거리를 구했다. 이렇게 구한 각 정점들에서, 해당 정점이 맥세권과 스세권을 만족하며 두 거리의 합이 최소인 값을 찾아 결과로 출력하는 방식을 택했다. 이 방식은 맥도날드와 스타벅스에 대한 총 두 번의 다익스트라를 사용하여 처음에 생각했던 방식보다 시간을 훨씬 절약할 수 있다.

정점의 개수와 도로의 개수를 각각 n,m에 저장하고, 각 도로에 대한 정보를 담을 adj라는 빈 리스트를 선언한다. 결과를 저장할 result에는 inf의 값을 넣어두고, 반복문을 사용하여 m개만큼 도로에 대한 정보를 입력하고, 이를 adj에 저장하도록 한다. 여기에서 x지점에서 y로 가는 정보만 저장한다면 y에서는 x로 가는 길을 찾지 못하기 때문에 x와 y에 대해 각각의 정보를 저장시킨다. mac에서는 맥도날드의 개수, 맥세권의 범위를 저장하고, maw에는 맥도날드의 위치를 저장한다. 마지막으로 sta에는 스타벅스의 개수, 스세권의 범위를 저장하고, stw에는 스타벅스의 위치를 저장한다. 먼저 맥도날드에 대해 다익스트라를 하기 위해서, mhq라는 리스트와 모든 값이 inf인 mcost라는 리스트를 만들어둔다. mcost에서 maw에 저장된 위치값을 모두 0으로 바꾸고, mhq에 해당 mcost값과 위치를 배열로 묶어 넣어둔다. 그리고 while문을 사용하여 mhq에서 가장 앞쪽의 값을 pop시키고, 해당 위치의 adj 요소들을 하나씩 확인하며 현재 쓰여있는 mcost의 값보다 두 가중치를 더한 값이 작을 경우 mcost의 값을 두 가중치의 합으로 변경한 뒤, heappush로 mhq에 mcost값과 해당 위치를 넣는다. 스타벅스에 대한 다익스트라도 위와 같은 과정으로 진행시킨 뒤, 반복문을 사용하여 맥도날드와 스타벅스가 없는 정점에 대해 맥세권, 스세권의 여부와 현재 result보다 작은지 여부를 확인하여 result의 값을 정한다. 모든 과정이 끝나면 result의 값을 출력시킨다. 이 과정을 코드로 나타내면 아래와 같다.

 

코드

from heapq import *
input=__import__('sys').stdin.readline
n,m=map(int,input().split())
adj=[[] for _ in range(n+1)]
result=float('inf')
for i in range(m):
  x,y,t=map(int,input().split())
  adj[x].append([y,t])
  adj[y].append([x,t])
mac=[*map(int,input().split())]
maw=[*map(int,input().split())]
sta=[*map(int,input().split())]
stw=[*map(int,input().split())]
mcost=[float('inf')]*(n+1)
mhq=[]
for i in range(mac[0]):
  mcost[maw[i]]=0
  heappush(mhq,[mcost[maw[i]],maw[i]])
while mhq:
  t,x=heappop(mhq)
  if mcost[x]!=t:
    continue
  for nx,nt in adj[x]:
    if mcost[nx]>t+nt:
      mcost[nx]=t+nt
      heappush(mhq,[mcost[nx],nx])
scost=[float('inf')]*(n+1)
shq=[]
for i in range(sta[0]):
  scost[stw[i]]=0
  heappush(shq,[scost[stw[i]],stw[i]])
while shq:
  t,x=heappop(shq)
  if scost[x]!=t:
    continue
  for nx,nt in adj[x]:
    if scost[nx]>t+nt:
      scost[nx]=t+nt
      heappush(shq,[scost[nx],nx])
for i in range(1,n+1):
  if i not in maw and i not in stw:
    if mcost[i]<=mac[1] and scost[i]<=sta[1] and mcost[i]+scost[i]<result:
      result=mcost[i]+scost[i]
if result!=float('inf'):
  print(result)
else:
  print(-1)

 

결과