[Python/BOJ] 14618. 총깡 총깡

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

14618번: 총깡 총깡

문제

동물 애호가 진서는 총깡총깡 뛰는 동물과 짝폴짝폴 뛰는 동물들을 K마리씩 키운다. 타지로 취업하게 된 진서는 내일 이사를 한다.

이사하게 될 집에서 같이 살게 될 룸메이트 일호는 동물을 싫어하기 때문에 진서는 근처의 집에 동물들을 한마리씩 맡길 예정이다.

진서가 동물들을 맡길 수 있는 집의 종류는 A형 집과 B형 집 2종류 이다.

우연하게도 짝폴짝폴 뛰는 동물과 총깡총깡 뛰는 동물, A형 집, B형 집의 수는 모두 같다.

진서는 총깡총깡 뛰는 동물들과 짝폴짝폴 뛰는 동물들을 같은 종류의 집에 통일 시켜 맡기고 싶다.

하지만 진서는 총깡총깡 뛰는 동물들을 약간 더 좋아하므로 각 집에서 동시에 출발하여 진서네 집으로 가장 빨리 도착하는 동물이 총깡총깡 뛰는 동물이길 원한다.

진서가 살게 될 집, A형 집, B형 집, A형 집도 B형 집도 아닌 집이 있는 지도가 주어질 때 총깡총깡 뛰는 동물이 A형 집에 살아야 할 지 B형집에 살아야 할지 출력하고 가장 빨리 도착하는 총깡총깡 뛰는 동물이 진서네 집으로 부터 얼마만큼 떨어져 있는지 출력하라.

(만약 총깡총깡 뛰는 동물들이 A형집에 살던 B형집에 살던 상관이 없는 경우는 A형집에 살기로 한다.)

입력

입력의 첫 번째 줄에 전체 집의 수 N과 집과 집사이를 연결하는 도로 M이 공백으로 주어진다. (3 ≤ N ≤ 5,000, 3 ≤ M ≤ 20,000)

입력의 둘째 줄에 진서의 집 J가 주어진다 (1 ≤ J ≤ N)

입력의 셋째 줄에 종류별 동물의 수 K가 주어진다. (2*K ≤ N)

입력의 넷째 줄에 K개의 A형 집이 공백으로 구분되어 주어진다.

입력의 다섯 번째 줄에 K개의 B형 집이 공백으로 구분되어 주어진다.

이후 M개의 줄에 X Y Z(1 ≤ X, Y ≤ N, 1 ≤ Z ≤ 100)가 주어진다. 이는 X번 집과 Y번 집 사이에 Z의 길이를 가지는 도로가 존재한다는 것이다.

출력

총깡총깡 뛰는 동물이 살게 될 집의 종류를 말한 뒤 다음줄에 거리를 출력한다.

A형 집에서만 진서의 집에 갈 수 있는 경우 A를 출력한 뒤 다음 줄에 거리를 출력, B형 집에서만 진서의 집에 갈 수 있는 경우 B를 출력한 뒤 다음 줄에 거리를 출력, A형 집, B형 집 둘다 진서의 집에 갈 수 없는 경우에는 –1을 출력한다.

풀이

Code

import sys
from heapq import heappush, heappop
input = sys.stdin.readline

# 1. 다익스트라 함수 정의
def djikstra(start, n) :
    node = [[] for _ in range(n+1)]
    # 1-1. 간선 정보 입력
    for _ in range(m) :
        a, b, c = map(int, input().split())
        node[a].append((b, c))
        node[b].append((a, c))
    # 1-2. 거리 정보 리스트 생성
    INF = float('INF')
    distance = [INF for _ in range(n+1)]
    # 1-3. 힙 리스트 생성 후 (시작 노드, 거리) 정보 입력
    heap = []
    heappush(heap, (start, 0))
    # 1-4. 거리 정보 리스트에 진서네 집 거리 설정
    distance[start] = 0
    # 1-5.
    while heap :
        # 1-5-1. 노드, 거리 정보 반환
        now, d = heappop(heap)
        # 1-5-2. 이미 처리된 경우 continue
        if distance[now] < d : continue
        # 1-5-3.
        for nn, c in node[now] :
            # 거리 계산
            cost = d + c
            # 계산된 거리가 현재 거리 리스트의 거리보다 짧을 경우
            if distance[nn] > cost :
                # 거리 정보 업데이트
                distance[nn] = cost
                # 힙에 (다음 노드, 거리) 정보 입력
                heappush(heap, (nn, cost))
    # 1-6. 거리 정보 리스트 반환
    return distance
def solution(n, m, j, k, home_A, home_B) :
    # 2. 다익스트라 함수 실행
    distance = djikstra(j, n)
    # 3. A형, B형 모두 진서네 집에 갈 수 없는 경우
    INF = float('INF')
    home_A_dist = [distance[idx] for idx in home_A]
    home_B_dist = [distance[idx] for idx in home_B]
    if home_A_dist.count(INF) + home_B_dist.count(INF) == 2 * k : print(-1)
    # 4. 이외의 경우
    else :
        # A와 B의 최소 거리가 같을 경우
        # A의 최소 거리가 더 가까울 경우
        if min(home_A_dist) <= min(home_B_dist) : print(f'A\n{min(home_A_dist)}')
        # B의 최소 거리가 더 가까울 경우
        else : print(f'B\n{min(home_B_dist)}')

if __name__ == "__main__" :
    n, m = map(int, input().split())
    j, k = int(input()), int(input())
    home_A, home_B = list(map(int, input().split())), list(map(int, input().split()))
    solution(n, m, j, k, home_A, home_B)