[LeetCode] 3123. Find Edges in Shortest Paths (Python)

안녕하세요!

오늘은 leetcode 3123번 Find Edges in Shortest paths 문제를 풀어보겠습니다. 

0~n-1 로 넘버링되어있는 n개의 노드를 가진 그래프가 있습니다. 이 그래프는 2d 배열인 'edges'로 대표되는 m edges로 이루어져 있는데요, edges[i] = [ai, bi, wi] 로 표현될 수 있고, 이는 노드 ai, bi간의 가중치 wi를 의미합니다. 0에서 n-1 노드 까지의 모든 최단 경로 중에서, edges[i] 가 포함되면 True, 그렇지 않으면 False를 출력하는 함수를 작성하는 문제입니다. 모든 결과를 'answer' 배열에 저장하고 출력해주면 됩니다. 

이 문제도 이전 시간에 풀었던 문제와 마찬가지로 Dijkstra 알고리즘을 활용해서 문제를 풀어보겠습니다. 

시작 노드에서 모든 노드까지의 최단 거리를 먼저 계산해줍니다. 이 과정에서 'parents'에 최단 경로 상의 노드를 저장해줍니다. 또 최단 경로에 포함된 edges를 stack이라는 리스트에 저장해줍니다. 마지막으로 real_edges를 통해 모든 엣지에 대해 최단 경로에 포함되는지 여부를 확인하고 answer에 T/F 결과를 추가해줍니다. 

import heapq

class Solution(object):
    def findAnswer(self, n, edges):
        """
        :type n: int
        :type edges: List[List[int]]
        :rtype: List[bool]
        """
        # 그래프 표현
        graph = {i: [] for i in range(n)}
        for a, b, w in edges:
            graph[a].append((b, w))
            graph[b].append((a, w)) # undirected

        # 부모 노드 추적 - 각 노드로 가는 최단 경로의 부모 노드 저장
        parents = {i: [] for i in range(n)}

        # 최단 거리 계산
        distance = [float('inf')] * n
        distance[0] = 0 # 시작 노드
        q = [(0, 0)] # 거리, 노드

        while q:
            dist, node = heapq.heappop(q)
            # 현재 탐색할 노드까지 거리가 기존에 저장된 거리보다 크면 스킵
            if dist > distance[node]:
                continue
            
            # 인접 노드까지의 거리가 기존에 저장된 거리보다 짧으면 업데이트
            for neighbor, weight in graph[node]:
                new_dist = dist + weight
                if new_dist < distance[neighbor]:
                    distance[neighbor] = new_dist
                    parents[neighbor] = [node] # 새로운 최단 경로로 업데이트
                    heapq.heappush(q, (new_dist, neighbor))
                elif new_dist == distance[neighbor]:
                    parents[neighbor].append(node) # 최단 거리가 여러 개일 경우 고려

        # 최단 경로에 포함된 edges
        # 최단 경로라면 경로가 중복되면 안되므로 이를 방지하고,
        # 포함 여부를 빠르게 확인하기 위해 집합 활용
        real_edges = set()
        stack = [n-1] # 목표 노드에서 최단 경로 역추적
        visited = set() # 방문한 노드 처리

        while stack:
            node = stack.pop() # 스택에서 노드를 꺼내 처리
            if node in visited: # 이미 방문한 노드는 스킵
                continue
            visited.add(node)
            for parent in parents[node]:    # 최단 경로 부모 노드들 확인 후 추가
                real_edges.add((parent, node))
                real_edges.add((node, parent))
                if parent not in visited:
                    stack.append(parent)

        # answer
        answer = []
        for a, b, w in edges:
            answer.append((a, b) in real_edges or (b, a) in real_edges)

        return answer

 

매우 어렵다!