[Algorithm][Softeer] 지우는 소수를 좋아해 (난이도 4) Python

---

---

풀이

처음 DFS로 풀었는데, 실패해서 다익스트라 알고리즘을 사용해 풀었다.

코드

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

def is_prime_number(x):
    if x <= 1:
        return False
    else:
        for i in range(2, int(math.sqrt(x))+1):
            if x % i == 0:
                return False
        return True

def dilkstra(start):
    q = []
    heapq.heappush(q, (0, start))
    distance[start] = 0
    while q:
        dist, now = heapq.heappop(q)
        if distance[now] < dist:
            continue
        for i in gym[now]:
            cost = max(dist, i[1])
            if cost < distance[i[0]]:
                distance[i[0]] = cost
                heapq.heappush(q, (cost, i[0]))

n, m = map(int, input().split())
gym = [[] for _ in range(n+1)]
distance = [INF] * (n+1)

for _ in range(m):
    a, b, c = map(int, input().split())
    gym[a].append((b,c))
    gym[b].append([a,c])

dilkstra(1)
x = distance[n]+1
while True:
    if is_prime_number(x):
        print(x)
        break
    x += 1

아래는 a,b,c 입력받아서 gym에 넣는 코드인데,
같은 경로가 여러 개이기 때문에, 최소 cost만 리스트에 추가하려고 했다.

그런데, 어차피 heapq 으로 최소값부터 돌기 때문에, 따로 처리 안해줘도 된다.

for _ in range(m):
    a, b, c = map(int, input().split())
    # 같은 경로에 여러 길 가능 -> 최소 레벨 넣기
    isIn = False
    for item in gym[a]:
        if item[0] == b:
            c = min(c, item[1])
            item[1] = c
            isIn = True
            break
    for item in gym[b]:
        if item[0] == a:
            c = min(c, item[1])
            item[1] = c
            isIn = True
            break
    if not isIn:
        if a != 1:
            gym[a].append([b,c])
            gym[b].append([a,c])
        else:
            gym[a].append([b,c])

## 아래 코드로 사용
for _ in range(m):
    a, b, c = map(int, input().split())
    gym[a].append((b,c))
    gym[b].append([a,c])