[백준 / 2644] 촌수계산 - Python

해당 문제는 다익스트라(Dijkstra) 알고리즘을 응용하여 풀었으므로 해당 알고리즘에 대한 설명은 아래 링크에서 확인하면 된다.

 

2021.03.18 - [# Coding/# Algorithm] - [보기 쉬운 Algorithm] 다익스트라(Dijkstra) 알고리즘 - Python

[보기 쉬운 Algorithm] 다익스트라(Dijkstra) 알고리즘 - Python

 

그래프를 먼저 입력받아 양방향 그래프를 작성한 후 모든 거리(가중치)를 1로 계산하여 다익스트라 알고리즘을 사용하였다.

<입력 1>

n = int(input())

graph = {i: [] for i in range(1, n+1)}

 

총인원수 n을 입력 받고, 그 크기만큼의 graph 키를 설정하였다.

<입력 2>

p1, p2 = map(int, input().split())

for _ in range(int(input())):
    x, y = map(int, input().split())
    graph[x].append(y)
    graph[y].append(x)

 

촌 수를 구하고 싶은 두 사람 p1, p2를 입력받는다.

 

관계의 개수를 입력받은 만큼 반복한다. 관계 선이 있는 두 사람 x, y를 입력받고 양방향으로 그래프를 그린다. 부모-자식, 자식-부모로 타고 갈 수 있기 때문이다.

<다익스트라 알고리즘>

import heapq


def dijkstra(graph, start_node):
    distances = {i: float('inf') for i in graph}
    distances[start_node] = 0

    queue = [[distances[start_node], start_node]]

    while queue:
        current_distance, current_node = heapq.heappop(queue)

        if current_distance > distances[current_node]:
            continue

        for next_node in graph[current_node]:
            distance = current_distance+1
            if distance < distances[next_node]:
                distances[next_node] = distance
                heapq.heappush(queue, [distance, next_node])

    return distances

 

속도 개선을 위해 최소 큐를 구현할 수 있게 heapq 모듈을 import했다.

 

start_node로 부터 계산되는 거리를 distances에 받을 수 있게 모든 값을 최대 수인 'inf'로 설정했다.

start_node의 거리는 포함하지 않으므로 0으로 설정한다.

 

queue에 우선순위가 되는 distances를 설정하고, 출발 노드인 start_node를 넣는다.

 

queue에서 현재 거리와 현재 기준 노드를 받아온다. 만약 이 거리가 distances에 있는 거리보다 길 경우 해당 값을 무시하고 다음 queue의 값으로 넘어간다.

 

짧을 경우 해당 노드로부터 갈 수 있는 모든 노드를 graph에서 받아온다. 모든 움직임의 거리가 1이므로, 새 거리는 현재 거리에서 1을 더한 값이고, 이 새거리가 distances에 있는 거리보다 짧은 경우, distances의 값을 갱신해주고, 해당 거리와 노드를 queue에 넣는다.

 

모든 반복이 끝난 후 distances를 반환해준다.

<호출 및 출력>

answer = dijkstra(graph, p1)[p2]

print(answer if answer != float('inf') else -1)

 

p1을 기준으로 반환된 distacnes의 p2와의 관계를 answer에 담는다.

만약 answer의 값이 'inf'인 경우 도착할 수 없으므로 -1로 값을 바꿔 출력하고, 아닌 경우 그대로 출력한다.

 

<전체 코드>

import heapq


def dijkstra(graph, start_node):
    distances = {i: float('inf') for i in graph}
    distances[start_node] = 0

    queue = [[distances[start_node], start_node]]

    while queue:
        current_distance, current_node = heapq.heappop(queue)

        if current_distance > distances[current_node]:
            continue

        for next_node in graph[current_node]:
            distance = current_distance+1
            if distance < distances[next_node]:
                distances[next_node] = distance
                heapq.heappush(queue, [distance, next_node])

    return distances


n = int(input())

graph = {i: [] for i in range(1, n+1)}

p1, p2 = map(int, input().split())

for _ in range(int(input())):
    x, y = map(int, input().split())
    graph[x].append(y)
    graph[y].append(x)

answer = dijkstra(graph, p1)[p2]

print(answer if answer != float('inf') else -1)