Problem statement

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

Solution

We can use bfs here, where state is (node, color of the last taken edge). Then we start with nodes (0, 1) and (0, 0) and perform bfs.

Complexity

It is O(E) for time and O(n) for space.

Code

class Solution:
    def shortestAlternatingPaths(self, n, R, B):
        G = defaultdict(list)
        for x, y in R:
            G[x] += [(y, 0)]
        for x, y in B:
            G[x] += [(y, 1)]
            
        queue = deque([(0, 0, 1), (0, 0, 0)])
        V = set([(0, 1), (0, 0)])
        ans = [float("inf")] * n
        
        while queue:
            dist, node, last = queue.popleft()
            ans[node] = min(ans[node], dist)
            for neib, col in G[node]:
                if (neib, col) not in V and col != last:
                    V.add((neib, col))
                    queue += [(dist + 1, neib, col)]
        
        return [i if i != float("inf") else -1 for i in ans]