Skip to the content.

leetcode [1129] 颜色交替的最短路径


Contact me:

Blog -> https://cugtyt.github.io/blog/index
Email -> cugtyt@qq.com
GitHub -> Cugtyt@GitHub


在一个有向图中,节点分别标记为 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]

来自题解:

记录红和蓝边的出度,如示例5:

红色的出度为[[1,2],[],[]],蓝色的出度为[[],[0],[]],

结果设为

[[0, 0], [INF, INF],[INF, INF]]

对于每种颜色的出度,从0开始,对于红色出度1,计算蓝色节点1的结果为

min(当前蓝色1节点,红色0到1 + 1)

如果蓝色节点有出度,继续到红色循环,否则,继续遍历红色的下一个出度2。以此类推。

class Solution:
    def shortestAlternatingPaths(self, n: int, red_edges: List[List[int]], blue_edges: List[List[int]]) -> List[int]:
        INF = 1e8
        rg = [[] for _ in range(n)]
        bg = [[] for _ in range(n)]
        for e in red_edges:
            rg[e[0]].append(e[1])
        for e in blue_edges:
            bg[e[0]].append(e[1])
        g = [rg, bg]
        res = [[INF, INF] for _ in range(n)]
        res[0] = [0, 0]
        self.dfs(g, 0, 0, res)
        self.dfs(g, 1, 0, res)
        out = [0] * n
        for i in range(n):
            out[i] = min(res[i][0], res[i][1])
            if out[i] == INF:
                out[i] = -1
        
        return out

    def dfs(self, g, col, i, res):
        other = 0 if col == 1 else 1
        for j in g[col][i]:
            if res[i][col] + 1 < res[j][other]:
                res[j][other] = res[i][col] + 1
                self.dfs(g, other, j, res)