[Softeer] S581 지우는 소수를 좋아해

[Softeer]

Question

Language: Python

해당 문제는 마지막 노드까지 도달하기 위한 최소 경로(거리)를 구하는 유형의 문제로, Dijkstra 알고리즘을 활용해서 풀이가 가능하다. 또한, 마지막 노드에 도달하기 위해 필요한 최소 레벨을 구할때 이 레벨은 소수이여야 하므로, 소수인지 여부를 판단하여, 소수가 아니면 소수를 찾을 때까지 값을 늘려주는 함수의 구현이 필요하다.

에라토스테네스 체를 활용하여 범위 내 저장된 소수를 모두 찾는 방법도 있지만, 해당 문제의 경우 수의 범위가 크기 때문에 제한된다.

소수 찾기

#소수 인지 여부 판별
def isPrime(number):
    for i in range(2,int(sqrt(number))+1):
        if number % i ==0:
            return False
    return True

#다음으로 큰 소수 찾기기
def find_next_prime(number):
    while not isPrime(number):
        number+=1
    return number

Solution

import sys
from heapq import heappop,heappush
from math import inf,sqrt

#소수 인지 여부 판별
def isPrime(number):
    for i in range(2,int(sqrt(number))+1):
        if number % i ==0:
            return False
    return True

#다음으로 큰 소수 찾기기
def find_next_prime(number):
    while not isPrime(number):
        number+=1
    
    return number
            
#마지막 노드까지 도달하는 최소 경로 찾기기
def solution():
    
    heap=[(0,1)]

    while heap:
        current_level,current_vertex=heappop(heap)

        #기존에 저장된 최소 레벨보다 큰 경우 넘어간다.
        if current_level > distances[current_vertex]:
            continue
        #최소 레벨 저장장
        distances[current_vertex]=current_level
        
        #다음 경로 탐색색
        for adj_vertex,next_level in graph[current_vertex]:
            next_level=max(current_level,next_level+1)

            if next_level < distances[adj_vertex]:
                distances[adj_vertex]=next_level
                heappush(heap,(next_level,adj_vertex))


if __name__ == "__main__":
    n_vertex,n_edge=map(int,input().split())
    graph=[[] for _ in range(n_vertex+1)]

    for _ in range(n_edge):
        v1,v2,weight=map(int,input().split())
        graph[v1].append((v2,weight))
        graph[v2].append((v1,weight))

    distances=[inf] *(n_vertex+1)

    solution()
    print(find_next_prime(distances[n_vertex]))