백준 18352번: 특정 거리의 도시 찾기

각 도시별 최단 노드를 구해 그 거리가K인 도시를 출력하면 된다.

노드별 최단 거리 이므로 다익스트라 알고리즘을 이용하면 된다.

 

import heapq

N, M, K, X = map(int, input().split())
graph = [[] for _ in range(N + 1)]
INF = 1e8
distance = [INF] * (N + 1)

for i in range(M):
    A, B = map(int, input().split())
    graph[A].append((B, 1))


def dijkstra(graph, start):
    q = []
    heapq.heappush(q, (0, start))
    distance[start] = 0

    while q:
        dist, now = heapq.heappop(q)

        if distance[now] < dist:
            continue

        for i in graph[now]:
            if dist + i[1] < distance[i[0]]:
                distance[i[0]] = dist + i[1]
                heapq.heappush(q, (dist + i[1], i[0]))


dijkstra(graph, X)
c=0
m=-1
for i in distance:
    c=c+1
    if i==K:
        m=0
        print(c-1)

if m==-1:
    print(-1)

 

각 도시의 갯수(노드 갯수), 각 도로 갯수(간선 갯수), 시작노드, 원하는 최단 거리를 입력받는다.

INF=1e8로 초기화해준다.

distance = [INF] * (N + 1), graph = [[] for _ in range(N + 1)] 로 초기화해준다.

N+1인 이유는 각 노드별로 바로 구할수 있게 하게끔이다.

예로, distance[4]=INF,graph[4]=[]  

각 가중치(도로의 길이)는 모두 1이므로

graph[도시].append(이동할 도시,거리)를 해주었다.

이제 우선순위큐를 이용해야한다.

우선 시작노드는 거리가 0이므로, 0과 시작노드인 x를 큐에 넣어준다.

이제 distance[x]=0으로 바뀌고, q에 요소가 없을때까지 반복문을 돌려준다.

q에 요소를 빼면 최단 거리와 노드가 나오는데, 이미 distance[노드]에 큐에 나온 최단거리보다 작은 거리가 있다면 밑에 반복문을 안돌려줘도된다.

최단 거리로만 이동해야하기때문이다!

이제 for문을 돌려 노드가 가리키는 노드와 가중치를 탐색한다.

거리가 distance[노드]보다 작으면, 값을 갱신해주고 큐에다 노드와 거리를 추가해준다.

이제 우선 순위 큐에서 pop을 하게 된다면 가장 거리가 작은 노드를 pop해줄수 있다

그 거리와 노드를 이용해 위에 동작을 반복해주고, q에 요소가 없을때까지 해주면된다.

이제 각 거리가 나온 distance를 return해주고, 거리가 k인 distance[노드]의 노드를 출력해주면 된다.