Python 中交替颜色的最短路径

pythonserver side programmingprogramming更新于 2025/11/12 22:36:17

假设我们有一个有向图，节点标记为 0、1、...、n-1。在此图中，每条边都用红色或蓝色着色，并且可能有自边或平行边。red_edges 中的每个 [i, j] 表示从节点 i 到节点 j 的红色有向边。类似地，blue_edges 中的每个 [i, j] 表示从节点 i 到节点 j 的蓝色有向边。我们必须找到一个长度为 n 的数组 answer，其中每个 answer[X] 是从节点 0 到节点 X 的最短路径的长度，使得边颜色沿路径交替(如果不存在这样的路径，则为 -1)。

因此，如果输入为 n = 3、red_edges = [[0,1],[1,2]] 和 blue_edges = []，则输出将为 [0, 1, -1]

为了解决这个问题，我们将遵循以下步骤 −

定义一个名为 bfs() 的方法，它将采用 re、be、f 和 n
定义一个名为 accessed 的集合，定义一个队列并插入一个三元组 [0, f, 0]
当 q 不为空时
- 设置三元组current、color、step 作为 q[0]，并从 q 中删除
- color := 补充 color 的值(真为假，反之亦然)
- res[current] := res[current] 和 step 的最小值
- 如果 color 非零，则
  - 对于 re[current] 中的每个 i
    - 如果对 (i, color) 不在访问中，则将 (i, color) 插入访问中并将 [i, color, step + 1] 插入 q
- 否则，当 color 为 0 时，则
  - 对于 be[current] 中的每个 i
    - 如果对 (i, color) 不在访问中，则将 (i, color) 插入访问中并将 [i, color, step + 1] 插入 q
在主方法 −
res := 一个无穷大值数组，其大小为 n
re 和 be := n 个空数组的数组
对于 r 中的每个元素 i:将 i[1] 插入 re[i[0]]
对于 b 中的每个元素 i:将 i[1] 插入 be[i[0]]
调用 bfs(re, be, false, n) 并调用 bfs(re, be, true, n)
对于 i，范围从 0 到 res 的长度 – 1
- 如果 res[i] = inf，则 res[i] := -1
返回 res

示例 (Python)

让我们看看下面的实现以便更好地理解 −

class Solution(object):
   def shortestAlternatingPaths(self, n, r, b):
      self.res = [float("inf")] * n
      re = [[] for i in range(n) ]
      be = [[] for i in range(n) ]
      for i in r:
         re[i[0]].append(i[1])
      for i in b:
         be[i[0]].append(i[1])
      self.bfs(re,be,False,n)
      self.bfs(re,be,True,n)
      for i in range(len(self.res)):
         if self.res[i] == float('inf'):
            self.res[i]=-1
      return self.res
   def bfs(self,re,be,f,n):
      visited = set()
      queue = [[0,f,0]]
      while queue:
         current,color,step = queue[0]
         queue.pop(0)
         color = not color
         self.res[current] = min(self.res[current],step)
         if color:
            for i in re[current]:
               if (i,color) not in visited:
                  visited.add((i,color))
                  queue.append([i,color,step+1])
         elif not color:
            for i in be[current]:
               if (i,color) not in visited:
                  visited.add((i,color))
                  queue.append([i,color,step+1])
ob = Solution()
print(ob.shortestAlternatingPaths(3, [[0,1], [1,2]], []))