[Python/BOJ] 14431. 소수마을

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

14431번: 소수마을

문제

소수 마을들의 주민들은 매우 특이한 규칙을 준수한다. 규칙은 바로 “가고 싶은 위치까지의 거리가 소수일 경우에만 간다”라는 것이다. 소수 마을의 주민 승욱이는 소수 마을에서 멀리 떨어진 A마을에 볼일이 있어 그곳까지 가야한다. 소수 마을에서 A마을까지의 단숨에 가고 싶지만 안타깝게도 두 마을의 거리는 소수가 아닐 경우에는 그럴 수가 없다. 그럴 경우에는 다른 마을들을 경유하여 가야한다. (경유하는 마을도 현재 위치에서의 거리가 소수일 경우에만 갈 수 있다.) 소수 마을과 경유할 수 있는 마을들, 그리고 A마을의 위치가 좌표평면 상으로 주어질 때, 승욱이가 소수 마을의 규칙을 준수하여 A마을로 갈 수 있는 최단의 길을 찾는 것을 도와주자. 소수 판정을 위해 마을 간의 거리는 정수 부분만으로 취급한다. 예를 들어, 거리가 3.1415라면 이를 버림하여 3만 취급한다.

입력

첫 번째 줄에 소수 마을의 위치 (X1,Y1)와 A마을의 위치 (X2,Y2)가 입력된다. 두 번째 줄에 경유할 수 있는 마을의 개수 N (0 ≤ N ≤ 4000)가 입력된다. 세 번째 줄부터 N+2번째 줄까지 경유 할 수 있는 마을들의 위치 (X3,Y3)가 입력된다. 단, 각 마을들의 좌표는 절댓값이 3000을 넘지 않는 정수이다.

출력

소수 마을의 규칙을 준수하여 A마을까지 가는 방법 중 제일 짧은 거리로 갈 수 있는 길의 거리합을 출력한다. 만약 소수 마을의 규칙을 준수하여 갈 수 있는 방법이 없는 경우 -1을 출력한다.

풀이

Code

import sys
from math import sqrt
from heapq import heappush, heappop
from collections import defaultdict
input = sys.stdin.readline

def solution(x1, y1, x2, y2, n) :
    # 1. 다익스트라 함수 생성
    def dijkstra(x, y) :
        # 1. 힙 생성 후 현재 (거리, 위치) 정보 입력
        heap = []
        heappush(heap, (0, x, y))
        # 2.
        while heap :
            # 2-1. 거리, 위치 인덱스 반환
            d, x, y = heappop(heap)
            # 2-2. 이미 처리된 경우 continue
            if distance[(x, y)] < d : continue
            # 2-3.
            for i, j in array :
                # 현재 노드와 다음 노드 간 거리 계산
                sub_d = int(sqrt((x - i) ** 2 + (y - j) ** 2))
                # 소수일 경우
                if check[sub_d] :
                    dist = d + sub_d
                    # 거리가 더 적을 경우
                    if dist < distance[(i, j)] :
                        # 거리 정보 업데이트
                        distance[(i, j)] = dist
                        # 힙에 정보 삽입
                        heappush(heap, (dist, i, j))

    # 2. 가능한 최대 길이 구하기
    max_len = int(sqrt(2 * (6000 ** 2))) + 1
    # 3. 소수 판별 리스트 생성
    check = [True for _ in range(max_len+1)]
    check[0], check[1] = False, False
    # 4. 에라토스테네스의 체
    for i in range(2, int(sqrt(max_len)) + 1) :
        if check[i] :
            for j in range(i+i, max_len+1, i) :
                if check[j] : check[j] = False
    # 5. 거리 정보 딕셔너리 생성
    distance = defaultdict(int)
    array = [(x2, y2)]
    # 6. 노드 정보 입력
    distance[(x1, y1)], distance[(x2, y2)] = 0, int(1e9)
    for _ in range(n) :
        x, y = map(int, input().split())
        distance[(x, y)] = int(1e9)
        array.append((x, y))
    # 7. 다익스트라
    dijkstra(x1, y1)
    # 8. 결과 출력
    print(-1) if distance[(x2, y2)] == int(1e9) else print(distance[(x2, y2)])

if __name__ == "__main__":
    x1, y1, x2, y2 = map(int, input().split())
    n = int(input())
    solution(x1, y1, x2, y2, n)