[백준] 14431번 소수마을 ★★★

💡문제접근

• 다익스트라 알고리즘을 작성하는 과정에서 좌표의 개념을 어떻게 적용해야할지 고민이 많았고 이를 좌표 그대로 처리하는 과정에서 IndexError가 이 부분을 해결하기 위해 많은 고민을 했다.
• 좌표를 그대로 넣지 않고 좌표에 해당하는 인덱스로 처리를 하여 다익스트라 알고리즘에 적용할 수 있도록 했다.
• 다른 사람의 코드를 보지 않고 효율성도 좋고 가독성도 좋은 코드를 작성하는 것을 목표로 하다보니까 어려운 문제로 보이는 문제에서 한계에 직면하면 오랫동안 시간을 투자하다보니까 실제 코딩테스트에서 실제로도 이렇게 오래 걸리는 것은 아닐까 많이 걱정이 들었다.

💡코드(메모리 : 37496KB, 시간 : 5936ms)

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

# 에라토스테네스의 체로 소수를 거르기
arr = [True] * 10001
arr[0] = False
arr[1] = False
for i in range(2, int(10001**0.5)+1):
    if arr[i]:
        for j in range(2*i, 10001, i):
            arr[j] = False

# 시작 좌표, 도착 좌표
graph = []
X1, Y1, X2, Y2 = map(int, input().strip().split())
N = int(input())
graph.append([X1, Y1])
graph.append([X2, Y2])
# 경유 좌표
for _ in range(N):
    a, b = map(int, input().strip().split())
    graph.append([a, b])

# 에라토스테네스의 체에 의한 소수 판별 함수
def is_prime(x):
    if arr[x]:
        return True
    else:
        return False

# 다익스트라 함수(좌표의 개념 → 값을 어떻게 처리해야할지 많은 고민...)
def dijkstra(start):
	# 소수마을 + 경유마을 + A마을
    # N번의 loop를 통한 경유마을 + 소수마을과 A마을을 추가해 N+2로 설정
    distance = [INF] * (N+2)
    distance[0] = 0
    q = []
    # (거리, 인덱스)
    heapq.heappush(q, [0, 0])
    while q:
        dist, now = heapq.heappop(q)
        now_x, now_y = graph[now]
        if distance[now] < dist:
            continue
        for i in range(N+2):
            next_x, next_y = graph[i]
            temp = int(math.sqrt((now_x - next_x) ** 2 + (now_y - next_y) ** 2))
            if is_prime(temp) and distance[i] > dist + temp:
                distance[i] = dist + temp
                heapq.heappush(q, [dist + temp, i])
    return distance

res = dijkstra(0)
# distance[1] : A마을까지의 최단 경로
if res[1] == INF:
    print(-1)
else:
    print(res[1])

💡소요시간 : 1h 37m