코딩테스트 문제

[백준/Python] 특정한 최단 경로

왕초보코딩러 2024. 10. 9. 20:50
728x90

 

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

 


방법은

노드 1 -> 노드 v1 -> 노드 v2 -> 노드 n

노드 1 -> 노드 v2 -> 노드 v1 -> 노드 n

두 가지 중 최소 경로를 구하고, 만약 INF 보다 크거나 같다면 -1을 출력하면 되는 문제였습니다.

 

 

일단 이 문제를 플로이드 워셜로 푸려고 했습니다

(이유는.. 800 ** 3을 1억번보다 더 적게 나온다고 잘못 계산해서..)

import sys
input = sys.stdin.readline
INF = int(1e9)

n, e = map(int, input().split())


graph = [[INF]*(n+1) for _ in range(n+1)]
for i in range(1, n+1):
    graph[i][i] = 0
    
for _ in range(e):
    a, b, c = map(int, input().split())
    graph[a][b] = c
    graph[b][a] = c
    
for k in range(1, n+1):
    for i in range(1, n+1):
        for j in range(1, n+1):
            graph[i][j] = min(graph[i][k]+graph[k][j], graph[i][j])

v1, v2 = map(int, input().split())
res = min(graph[1][v1] + graph[v1][v2] + graph[v2][n], 
          graph[1][v2] + graph[v2][v1] + graph[v1][n]         
        )
if res >= INF :
    print(-1)
else:
    print(res)

 

물론 800 ** 3은 5억번이 넘어가기 때문에 당연히 시간 초과로 틀렸습니다

 

근데 제가 플로이드 워셜의 시간을 줄이는 방법을 답변 받았습니다.

1. 함수 내에서 수행하기

파이썬이 전역의 변수 접근보다 함수 내 변수 접근이 훨씬 빠르다고 합니다

 

2. 리스트 인덱스 접근 최소화

물론 리스트의 접근은 시간 복잡도가 O(1)이지만,

리스트가 매우 크고 캐시 메모리로 관리하기 힘든 경우 메모리 접근 속도가 느려질 수 있다고 합니다.

그래서 반복적으로 인덱스 접근을 한다면 시간이 느려질 수 있다는 것입니다.

그래서 graph[i][j]를 접근하는 대신, 변수에 graph[i]를 저장해두고 접근한다면 더 빨라지는 것입니다.

 

일단 그래서 플로이드 워셜 함수를 만들었습니다.

def fw(graph):
    for k in range(1, n+1):
        gk = graph[k]
        for i in range(1, n+1):
            gi = graph[i]
            x = gi[k]
            for j in range(1, n+1):
                if i == j:
                    continue
                gi[j] = min(x+gk[j], gi[j])

graph[i]와 graph[k]를 많이 접근하기 때문에 변수에 저장하고

graph[i][k] 또한 계속 접근하기 때문에 변수에 저장했습니다

 

또한 i와 j가 똑같을 때는(최솟값은 어차피 0) continue를 하여 최적화하였습니다.

 

 

이 방법을 사용하여 pypy3으로 통과 하였습니다

 


물론 다익스트라를 사용하여 풀 수 있습니다.

 

import sys
import heapq
input = sys.stdin.readline
INF = int(1e9)

def dijkstra(start, end, graph, distance):
    q = []
    heapq.heappush(q, (0, start))
    distance[start] = 0
    
    while q:
        cost, node = heapq.heappop(q)
        
        if cost > distance[node]:
            continue
        for i in graph[node]: # i[0]은 cost, i[1]은 노드
            c = cost + i[0]
            if c < distance[i[1]]:
                distance[i[1]] = c
                heapq.heappush(q, (c, i[1]))
    return distance[end] 

n, e = map(int, input().split())
g =  [[] for _ in range(n+1)]
for _ in range(e):
    a, b, c = map(int, input().split())
    g[a].append((c, b))
    g[b].append((c, a))
    
v1, v2 = map(int, input().split())

def min_distance(s, e):
    if s == e:
        return 0
    distance = [INF] * (n+1)
    return dijkstra(s, e, g, distance)

tov1 = min_distance(1, v1)
v1tov2 = min_distance(v1, v2)
v2ton = min_distance(v2, n)
tov2 = min_distance(1, v2)
v2tov1 = min_distance(v2, v1)
v1ton = min_distance(v1, n)

res = min(
        (tov1 + v1tov2 + v2ton), # 1 -> v1 -> v2 -> n
        (tov2 + v2tov1 + v1ton), # 1 -> v2 -> v1 -> n
    )
if res >= INF:
    print(-1)
else:
    print(res)

 

 

일단 함수 실행이 더 빠르다는 사실과 리스트 접근을 최소화하면 시간을 줄일 수 있다는 것을 알게 되어 재밌었습니다!!

'코딩테스트 문제' 카테고리의 다른 글

[Python/백준] 트리 순회  (1) 2024.10.15
[백준/Python] 트리의 부모 찾기  (0) 2024.10.14
[Python/백준] 운동  (1) 2024.10.05
[Python/백준] 구간 합 구하기 5  (1) 2024.09.28
[Python/백준] RGB 거리  (0) 2024.09.23