[백준][아직못풀었음] 18352 특정 거리 도시 찾기 Python

1. 문제

1번부터 N번까지의 도시와 M개의 단방향 도로가 존재할 때 (모든 도로의 거리는 1이다)
특정한 도시 X로부터 출발하여 최단 거리가 K인 모든 도시의 번호를 구하는 프로그램을 작성하시오
2 <= N <= 300,000, 1 <= M <= 1,000,000, 1 <= K <= 300,000, 1 <= X <= N 

 

2. 알고리즘 분류

* 그래프 이론

* 그래프 탐색

* 너비 우선 탐색

* 최단 경로

* 데이크스트라

 

3. 접근 방식

Week02 퀴즈 주제인 다익스트라를 이용하여 문제를 푸려고 한다

 

4. 전체 코드

첫번째 시도 

다익스트라 함수를 별개로 구현하고, N, M, K, X 를 입력 받아 딕셔너리 형태로 graph 를 만들었다 

그리고 딕셔너리에서 items() 로 node 와 edge 를 꺼내 edge 와 K 를 비교해서 

result list 에 append 하고, 나중에 .join() 으로 정리해서 print 했다

import heapq

def dijkstra(graph, start):
    # 거리 테이블을 무한대로 초기화하고, (노드, 거리) 형태로 저장
    dist = {node: float('inf') for node in graph}
    dist[start] = 0
    # 우선순위 큐 초기화 
    # 거리를 기준으로 정렬하기 위해 (거리, 노드) 형태로 저장 
    queue = [(0, start)]  

    while queue:
    	# heapq.heappop() 로 queue 에서 제일 짧은 거리를 가진 노드를 꺼낸다 
        current_distance, current_node = heapq.heappop(queue)

        # 현재 처리 중인 노드의 거리가 기록된 거리보다 크다면 
        # 이미 최적의 경로로 처리되었으므로 무시한다 
        # current_distance 는 우선순위 큐에서 꺼낸 값
        # dist[current_node] 현재 노드까지의 최단 거리
        if current_distance > dist[current_node]:
            continue

		# for 문으로 인접 노드와 가중치 차례로 확인 
        for next_cost, next_node in graph[current_node].items():
            new_distance = current_distance + next_node

            # 더 짧은 경로가 발견되면 거리 테이블을 갱신하고 우선순위 큐에 추가
            if dist[next_cost] > new_distance:
                dist[next_cost] = new_distance
                heapq.heappush(queue, (new_distance, next_cost))

    return dist

if __name__ == '__main__':
    # N 도시의 개수 M 도로의 개수 K 거리 정보 X 출발 도시의 번호 
    N, M, K, X = map(int, input().split())
    
    # K = int(input())

    # graph = {
    #     '1' : {'2': 1, '3': 1},
    #     '2' : {'3': 1, '4': 1},
    #     '3' : {},
    #     '4' : {}
    # }

    graph = {i: {} for i in range(1, N+1)}
    for _ in range(M):
        u, v = map(int, input().split())
        graph[u][v] = 1

    node_result = []
    result = dijkstra(graph, X)
    for node, edge in result.items():
        if edge == K:
            node_result.append(str(node))
    
    if not node_result:
        print(-1)
    else:
        print(" ".join(node_result))

근데 시간 초과가 떴다

ChatGPT 한테 물어보니,

1) 각 간선이 동일한 가중치를 가질 때 인접 리스트 형태로만 구현하여도 충분하니

딕셔너리 대신 리스트로, 간단한 방문 배열을 사용하는 것이 효율적이라고 했다

2) 큐에 너무 많은 노드가 들어가면 시간이 많이 소요될 수 있으니 가중치를 기록할 필요 없이 

(graph[u][v] = 1 부분) 단순히 인접 리스트에 추가하는 방식을 추천하고

3) input() 대신 sys.stdin.read 사용하여 입력을 빠르게 처리하는 것이 좋다고 했다 

 

일단은 다익스트라 개념으로 넘어가서 딕셔너리를 리스트로 바꾸는 것부터 해야겠다