leetcode-1129. Shortest Path with Alternating Colors

题目

You are given an integer n, the number of nodes in a directed graph where the nodes are labeled from 0 to n - 1. Each edge is red or blue in this graph, and there could be self-edges and parallel edges.

You are given two arrays redEdges and blueEdges where:

redEdges[i] = [ai, bi] indicates that there is a directed red edge from node ai to node bi in the graph, and
blueEdges[j] = [uj, vj] indicates that there is a directed blue edge from node uj to node vj in the graph.
Return an array answer of length n, where each answer[x] is the length of the shortest path from node 0 to node x such that the edge colors alternate along the path, or -1 if such a path does not exist.

Example 1:

Input: n = 3, redEdges = [[0,1],[1,2]], blueEdges = []
Output: [0,1,-1]

Example 2:

Input: n = 3, redEdges = [[0,1]], blueEdges = [[2,1]]
Output: [0,1,-1]

Constraints:

1 <= n <= 100
0 <= redEdges.length, blueEdges.length <= 400
redEdges[i].length == blueEdges[j].length == 2
0 <= ai, bi, uj, vj < n

解题思路

图的最短路径的变形题，状态里保存当前节点选择的下一条路径的颜色即可
时间复杂度是$o(V+E)$，其中V和E分别是节点数和边数
空间复杂度是$o(V)$

0和1互相更换，可以使用xor

代码

class Solution:
    def shortestAlternatingPaths(self, n: int, redEdges: List[List[int]], blueEdges: List[List[int]]) -> List[int]:
        import collections
        # build graph
        red_neighbors = {}
        blue_neighbors = {}
        for edge_src, edge_dst in redEdges:
            if edge_src not in red_neighbors:
                red_neighbors[edge_src] = []
            red_neighbors[edge_src].append(edge_dst)
        for edge_src, edge_dst in blueEdges:
            if edge_src not in blue_neighbors:
                blue_neighbors[edge_src] = []
            blue_neighbors[edge_src].append(edge_dst)
        # bfs
        res = [-1] * n
        neighbor_choices = [red_neighbors, blue_neighbors]
        # node_index, path_len, next_color
        queue = collections.deque([(0, 0, 0), (0, 0, 1)])
        visited_nodes = set()
        while queue:
            node, path, color = queue.popleft()
            if (node, color) in visited_nodes:
                continue
            visited_nodes.add((node, color))
            res[node] = min(path, res[node] if res[node] != -1 else path)
            for neighbor in neighbor_choices[color].get(node, []):
                queue.append((neighbor, path + 1, int(not(color))))
        return res