[Softeer]지우는 소수를 좋아해 파이썬 풀이

Softeer 지우는 소수를 좋아해¶

문제 page: https://softeer.ai/practice/info.do?eventIdx=1&psProblemId=582&sw_prbl_sbms_sn=64055

문제유형: 다익스트라

# 일반적인 다익스트라 알고리즘
if 현재 cost > 경유 cost:
    현재 cost = 경유 cost

다만 문제에서는 현재 cost를 누적해서 더해나가는 최단경로 문제가 아닌 경로상에서 주어진 cost 중 max값을 찾는 문제입니다.

In [3]:

from IPython.display import Image
Image('지우는소수를좋아해1.png')

Out[3]:

즉 비교를 할 때 costs[node] + now 가 아닌 max(costs[node], now)를 통해 값을 갱신해나가야 합니다.
또한 정답의 경우 문제 조건에 따라 N번 체육관에 도달하기 위한 레벨보다 높은 소수로 출력을 해줘야 합니다.

ref: https://www.youtube.com/watch?v=cswJ1h-How0&list=PLRx0vPvlEmdAghTr5mXQxGpHjWqSz0dgC&index=9 소수 판별 알고리즘
소수 판별 알고리즘
소수: 1과 자기 자신외에 나누어 떨어지지 않는 수
소수의 정의에 따라 알고리즘을 작성하면 됩니다.

def isPrime(n):
    if n==1:
        return True
    if n==2:
        return False
    for i in range(2, n):
        if n%i == 0:
            return False
    return True

Tip) 특정한 수 n은 $sqrt{n}$ 이하인 수와 $sqrt{n}$ 이상인 수의 곱으로 이루어져 있습니다.
따라서 n을 반으로 갈라 $sqrt{n}$이하 까지만 나누어 떨어지는지 검사를 해줘도 소수판별을 할 수 있습니다.
ex) 16 = 1 x 16, 2 x 8, 4 x 4

# speed up version
def isPrime(n):
    if n==1:
        return True
    if n==2:
        return False
    for i in range(2, int(sqrt(n))+1):
        if n%i == 0:
            return False
    return True

cf) 해당 문제의 경우 다익스트라 알고리즘을 heapq을 이용해서 짜지 않을 시에 시간 제한에 걸리게 됩니다.
시간 제한에 걸려서 sys.stdin.readline(), 소수판별 알고리즘 개선 등 다양한 방법을 사용했지만 결국 heapq를 사용하지 않은 것이 문제였습니다..

In [4]:

import math
import heapq
N, M = [int(x) for x in input().split()]

visited = [False]*(N+1) # Not use 0 index
costs = [math.inf]*(N+1) # Not use 0 index
graph = {}
for i in range(1, N+1):
    graph[i] = []
for _ in range(M):
    a, b, c = [int(x) for x in input().split()]
    graph[a].append( (b, c) )
    graph[b].append( (a, c) )
    
def min_cost_graph(s_node):
    costs[s_node] = 0
    # for node, cost in graph[s_node]:
    #     costs[node] = cost
    q = []
    heapq.heappush(q, (costs[s_node], s_node))
    while q:
        _, now = heapq.heappop(q)
        visited[now] = True

        for node, cost in graph[now]:
            if costs[node] > max(costs[now], cost):
                costs[node] = max(costs[now], cost)
                heapq.heappush(q, (costs[node], node))

min_cost_graph(1)
def isPrime(n):
    if n == 1:
        return True
    if n == 2:
        return False
    for i in range(2, int(math.sqrt(n))+1):
        if n%i==0:
            return False
    return True

costs[-1] += 1
while isPrime(costs[-1]) != True:
    costs[-1] += 1
print(costs[-1])

10 13
1 2 5
1 3 1
1 4 2
2 5 5
3 5 4
3 6 1
4 6 1
4 7 3
5 8 5
6 9 4
7 9 2
8 10 5
9 10 3
5

In [ ]:

# 시간 제한
import math
from collections import deque
import heapq
N, M = [int(x) for x in input().split()]

visited = [False]*(N+1) # Not use 0 index
costs = [math.inf]*(N+1) # Not use 0 index
graph = {}
for i in range(1, N+1):
    graph[i] = []
for _ in range(M):
    a, b, c = [int(x) for x in input().split()]
    graph[a].append( (b, c) )
    graph[b].append( (a, c) )
    
def min_cost_graph(s_node):
    costs[s_node] = 0
    # for node, cost in graph[s_node]:
    #     costs[node] = cost
    q = deque([s_node])
    while q:
        now = q.popleft()
        visited[now] = True

        for node, cost in graph[now]:
            if costs[node] > max(costs[now], cost):
                costs[node] = max(costs[now], cost)
        
        min_value = math.inf
        min_idx = None
        for i in range(1, N+1):
            if visited[i] == False:
                if min_value > costs[i]:
                    min_value = costs[i]
                    min_idx = i
        if min_idx is not None:
            q.append(min_idx)
            
min_cost_graph(1)
def isPrime(n):
    if n == 1:
        return True
    if n == 2:
        return False
    for i in range(2, int(math.sqrt(n))+1):
        if n%i==0:
            return False
    return True

costs[-1] += 1
while isPrime(costs[-1]) != True:
    costs[-1] += 1
print(costs[-1])