【Leetcode】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

来源：力扣（LeetCode）
链接：https://leetcode-cn.com/problems/shortest-path-with-alternating-colors

思路

BFS，从节点0起，每次走一步，这样保证到达每个节点时路径是最短的。因为题目要的是红蓝相间，所以“红蓝红”、“蓝红蓝”这两种分别以“红”开头和以“蓝”开头的都要考虑。比较巧妙的方法是利用队列，使得这两种依次走一步，正好能求出结果。

代码

class Solution:
    def shortestAlternatingPaths(self, n: int, red_edges: List[List[int]], blue_edges: List[List[int]]) -> List[int]:
        from collections import defaultdict

        red_colloction = defaultdict(list)
        blue_collection = defaultdict(list)
        for point in red_edges:
            red_colloction[point[0]].append(point[1])
        for point in blue_edges:
            blue_collection[point[0]].append(point[1])

        red_visit = {0:True}
        blue_visit = {0:True}
        queue = [(0, 0, "red"), (0, 0, "blue")]
        ans = [float("inf") for i in range(n)]
        while queue:
            curr, level, color = queue.pop(0)
            if(color == "red"):
                for point in red_colloction[curr]:
                    if(point not in red_visit):
                        red_visit[point] = True
                        queue.append((point, level+1, "blue"))
                        ans[point] = min(ans[point], level+1)
            else:
                for point in blue_collection[curr]:
                    if(point not in blue_visit):
                        blue_visit[point] = True
                        queue.append((point, level+1, "red"))
                        ans[point] = min(ans[point], level+1)
        ans = [-1 if x==float("inf") else x for x in ans]
        ans[0] = 0
        return ans