이해

• 1번부터 N번까지의 도시와 거리가 1인 M개의 단방향 도로가 존재한다.
• 특정한 도시 X로부터 출발하여 출발할 수 있는 모든 도시 중에서, 최단 거리가 정확히 K인 모든 도시들의 번호를 출력하라.

계획

• 다익스트라 알고리즘을 사용하여 도시 X로부터 다른 모든 도시까지의 최단 거리를 계산한다.
• 최단 거리가 정확히 K인 모든 도시들의 번호를 출력한다.

실행

import heapq

def dijkstra(start, graph, distance):
    distance[start] = 0
    q = []
    heapq.heappush(q, (0, start))
    while q:
        dist, now = heapq.heappop(q)
        if dist > distance[now]:
            continue
        for adj in graph[now]:
            cost = dist + 1
            if cost < distance[adj]: 
                distance[adj] = cost
                heapq.heappush(q, (cost, adj))


n, m, k, x = map(int, input().split())
graph = [[] for _ in range(n+1)]

for _ in range(m):
    a, b = map(int, input().split())
    graph[a].append(b)

INF = int(1e9)
distance = [INF]*(n+1)

dijkstra(x, graph, distance)

result = [i for i in range(1, n+1) if distance[i] == k]
if len(result) == 0:
    print(-1)
else:
    for i in result:
        print(i)

개선

• 최단경로를 찾는 문제에서 모든 간선의 weight가 동일한 경우 bfs 알고리즘을 사용하여 구현이 가능하다.
  • 다익스크라 알고리즘의 시간 복잡도는 O(|E| * log|E|)이지만(보통 |E| <= |V|^2 이므로, O(|E|log|V|)라고 볼 수도 있음) bfs 알고리즘의 시간 복잡도는 O(|E| + |V|)이므로 bfs 알고리즘을 사용하는 것이 훨씬 효율적이다.

from collections import deque

def bfs(start, graph, distance):
    q = deque([start])
    distance[start] = 0
    while q:
        now = q.popleft()
        for adj in graph[now]:
            if distance[adj] == INF:
                distance[adj] = distance[now] + 1
                q.append(adj)

n, m, k, x = map(int, input().split())
graph = [[] for _ in range(n+1)]

for _ in range(m):
    a, b = map(int, input().split())
    graph[a].append(b)

INF = int(1e9)
distance = [INF]*(n+1)


bfs(x, graph, distance)

result = [i for i in range(1, n+1) if distance[i] == k]
if len(result) == 0:
    print(-1)
else:
    for i in result:
        print(i)

회고

• 쉬웠다.
• 최단거리 구하는 문제에서 그래프에서 모든 간선의 가중치가 같을 때, bfs 알고리즘을 사용하여 최단 거리를 구할 수 있다는 것을 까먹었다. 꼭 기억해야겠다.