LeetCode 1129. Shortest Path with Alternating Colors

Graph Theory 문제 4번
https://leetcode.com/problems/shortest-path-with-alternating-colors/

나의 풀이(35ms, 5.03%)

• heap을 이용해서 풀었다.
• 모든 엣지의 가중치가 1이므로 굳이 heap을 사용하지 않고 단순한 BFS 풀이가 자연스럽다.

heap

class Solution:
    def shortestAlternatingPaths(self, n: int, redEdges: List[List[int]], blueEdges: List[List[int]]) -> List[int]:
        dist = collections.defaultdict(list) # (node, color) : dist
        graph = collections.defaultdict(list) # (node, color)
        heap = [(0, 0, 0, 'red'), (0, 0, 0, 'blue')] # dist, node_from, node_to, color
        swap = {'red':'blue', 'blue':'red'}
        result = [-1] * n

        for a, b in redEdges:
            graph[a].append((b, 'red'))
        for a, b in blueEdges:
            graph[a].append((b, 'blue'))

        dist[(0, 'red')] = 0
        dist[(0, 'blue')] = 0

        while heap:
            d, u, v, color = heapq.heappop(heap)
            if (u, swap[color]) in dist:
                if (v, color) not in dist:
                    dist[(v, color)] = dist[(u, swap[color])] + 1
                else:
                    dist[(v, color)] = min(dist[(v, color)], dist[(u, swap[color])] + 1)
                for w, new_color in graph[v]:
                    if (w, new_color) not in dist:
                        heapq.heappush(heap, (dist[(v, color)], v, w, new_color))

        for x in dist:
            result[x[0]] = min(result[x[0]], dist[x]) if result[x[0]] != -1 else dist[x]
 
        return result

다른 풀이1(1ms, 85.74%)

BFS

class Solution:
    def shortestAlternatingPaths(self, n: int, redEdges: List[List[int]], blueEdges: List[List[int]]) -> List[int]:
        graph = defaultdict(lambda: {'red': [], 'blue': []})
        for u, v in redEdges:
            graph[u]['red'].append(v)
        for u, v in blueEdges:
            graph[u]['blue'].append(v)

        result = [-1] * n
        visited = [[False, False] for _ in range(n)]  # visited[node][0: red, 1: blue]

        q = deque()
        q.append((0, 0, None))  # node, steps, prev_color

        while q:
            node, steps, color = q.popleft()
            if result[node] == -1:
                result[node] = steps
            # 다음으로 갈 색상은 현재 색상과 반대
            for next_color in ('red', 'blue'):
                if next_color == color:
                    continue
                color_idx = 0 if next_color == 'red' else 1
                for nei in graph[node][next_color]:
                    if not visited[nei][color_idx]:
                        visited[nei][color_idx] = True
                        q.append((nei, steps + 1, next_color))

        return result