Участник:Novitskiy97/Shortest Path with Alternating Colors

Материал из DISCOPAL
Перейти к: навигация, поиск

https://leetcode.com/problems/shortest-path-with-alternating-colors/

Python 3

class Solution:
    def shortestAlternatingPaths(self, n: int, red_edges: List[List[int]], blue_edges: List[List[int]]) -> List[int]:
 
        adj = defaultdict(list)
        output = [-1] * n
        output[0] = 0
 
        for v, u in red_edges:
            adj[v].append([u, "r"])
 
        for v, u in blue_edges:
            adj[v].append([u, "b"])
 
        bfs = deque()
        for v, col in adj[0]:
            bfs.append([v, col, 1])
 
        visited = set()
        while bfs:
            v, col, steps = bfs.popleft()
 
            if output[v]==-1:
                output[v] = steps
 
            if (v, col) not in visited:
                visited.add((v,col))
                for u, col2 in adj[v]:
                    if (col != col2) and ((u, col2) not in visited):
                        bfs.append([u, col2, steps + 1])
 
        return output