LeetCode-Python-1129. 颜色交替的最短路径

在一个有向图中，节点分别标记为 0, 1, ..., n-1。这个图中的每条边不是红色就是蓝色，且存在自环或平行边。

red_edges 中的每一个 [i, j] 对表示从节点 i 到节点 j 的红色有向边。类似地，blue_edges 中的每一个 [i, j] 对表示从节点 i到节点 j 的蓝色有向边。

返回长度为 n 的数组 answer，其中 answer[X] 是从节点 0 到节点 X 的最短路径的长度，且路径上红色边和蓝色边交替出现。如果不存在这样的路径，那么 answer[x] = -1。

 

示例 1：

输入：n = 3, red_edges = [[0,1],[1,2]], blue_edges = []
输出：[0,1,-1]

示例 2：

输入：n = 3, red_edges = [[0,1]], blue_edges = [[2,1]]
输出：[0,1,-1]

示例 3：

输入：n = 3, red_edges = [[1,0]], blue_edges = [[2,1]]
输出：[0,-1,-1]

示例 4：

输入：n = 3, red_edges = [[0,1]], blue_edges = [[1,2]]
输出：[0,1,2]

示例 5：

输入：n = 3, red_edges = [[0,1],[0,2]], blue_edges = [[1,0]]
输出：[0,1,1]

 

提示：

• 1 <= n <= 100
• red_edges.length <= 400
• blue_edges.length <= 400
• red_edges[i].length == blue_edges[i].length == 2
• 0 <= red_edges[i][j], blue_edges[i][j] < n

思路：

问最短路径就上BFS。

class Solution(object):
    def shortestAlternatingPaths(self, n, red_edges, blue_edges):
        """
        :type n: int
        :type red_edges: List[List[int]]
        :type blue_edges: List[List[int]]
        :rtype: List[int]
        """
        from collections import defaultdict, deque
        red_nei, blue_nei = defaultdict(set), defaultdict(set)
        
        for start, end in red_edges: #记录下所有的从红色路径可以到达的邻居节点
            red_nei[start].add(end)
        for start, end in blue_edges: #记录下所有的从蓝色路径可以到达的邻居节点
            blue_nei[start].add(end)
            
        visited = set()
        queue = deque()
        queue.append((0, -1)) # -1 无色， 0 红色， 1蓝色
        distance = 0
        res = [-1] * n
        while queue:
            next_queue = deque()
            for cur, color in queue:
                if res[cur] == -1: #BFS保障了第一次到达终点时的距离一定是最短距离
                    res[cur] = distance
                if color == -1 or color == 0: #当前是红色，下一次在蓝色里找
                    for nei in blue_nei[cur]:
                        if (cur, nei, 1) not in visited:
                            visited.add((cur, nei, 1))
                            next_queue.append((nei, 1))
                if color == -1 or color == 1: #当前是蓝色，下一次在红色里找
                    for nei in red_nei[cur]:
                        if (cur, nei, 0) not in visited:
                            visited.add((cur, nei, 0))
                            next_queue.append((nei, 0))
            queue = next_queue
            distance += 1
        return res