백준 21278 : 호석이 두마리 치킨

백준 21278 문제 보기

import sys
input = sys.stdin.readline
from itertools import combinations
import heapq
INF = sys.maxsize

# 입력
N, M = map(int, input().split())
roads = [ [0 if i == j else INF for j in range(N+1)] for i in range(N+1) ]
results = []
for i in range(M):
  a, b = map(int, input().split())
  roads[a][b], roads[b][a] = 1, 1

# Floyd Warshall
def fw():  
  for i in range(1, N+1):
    for s in range(1, N+1):
      for e in range(1, N+1):
        if s == e : continue
        roads[s][e] = min(roads[s][i]+roads[i][e], roads[s][e])

# 거리 구해주는 함수. 왕복이 아니고 편도.
def s(f1, f2):
  r = 0
  for i in range(1, N+1):
    if roads[f1][i] == INF or roads[f2][i] == INF: return INF
    r += min(roads[f1][i], roads[f2][i])
  return r

fw()

for c in combinations(range(1, N+1), 2):
  r = s(c[0], c[1])
  if r == INF: continue
  t = [min(c[0], c[1]), max(c[0], c[1]), r*2]
  heapq.heappush(results, ((t[2], t[0], t[1]), t))

print(*results[0][1])

처음에는 다익스트라로 접근했다.
문제 조건을 잘못 보고 정렬을 잘못해서 틀렸고, 정렬조건만 바꿔주면 맞을 줄 알았는데 시간초과가 났다.

sort 함수를 걷어내고 heap을 이용해 정렬하는 방식으로 바꾸면 시간초과가 안날줄 알았는데 그래도 시간 초과가 났다.

생각을 해보니 모든 경로로 가는 비용을 전부 구해야 하는 문제였고, 플로이드 와샬로 알고리즘을 바꿨다.

출력 초과는 인간미....