Skip to content

1129. Shortest Path with Alternating Colors 👍

  • Time: $O(n)$
  • Space: $O(n)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
enum class Color { kInit, kRed, kBlue };

class Solution {
 public:
  vector<int> shortestAlternatingPaths(int n, vector<vector<int>>& redEdges,
                                       vector<vector<int>>& blueEdges) {
    vector<int> ans(n, -1);
    vector<vector<pair<int, Color>>> graph(n);  // graph[u] := [(v, edgeColor)]
    queue<pair<int, Color>> q{{{0, Color::kInit}}};  // [(u, prevColor)]

    for (const vector<int>& edge : redEdges) {
      const int u = edge[0];
      const int v = edge[1];
      graph[u].emplace_back(v, Color::kRed);
    }

    for (const vector<int>& edge : blueEdges) {
      const int u = edge[0];
      const int v = edge[1];
      graph[u].emplace_back(v, Color::kBlue);
    }

    for (int step = 0; !q.empty(); ++step)
      for (int sz = q.size(); sz > 0; --sz) {
        const auto [u, prevColor] = q.front();
        q.pop();
        ans[u] = ans[u] == -1 ? step : ans[u];
        for (auto& [v, edgeColor] : graph[u]) {
          if (v == -1 || edgeColor == prevColor)
            continue;
          q.emplace(v, edgeColor);
          v = -1;  // Mark (u, v) as used.
        }
      }

    return ans;
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
enum Color { kInit, kRed, kBlue }

class Solution {
  public int[] shortestAlternatingPaths(int n, int[][] redEdges, int[][] blueEdges) {
    int[] ans = new int[n];
    Arrays.fill(ans, -1);
    // graph[u] := [(v, edgeColor)]
    List<Pair<Integer, Color>>[] graph = new List[n];
    // [(u, prevColor)]
    Queue<Pair<Integer, Color>> q = new ArrayDeque<>(List.of(new Pair<>(0, Color.kInit)));

    for (int i = 0; i < n; ++i)
      graph[i] = new ArrayList<>();

    for (int[] edge : redEdges) {
      final int u = edge[0];
      final int v = edge[1];
      graph[u].add(new Pair<>(v, Color.kRed));
    }

    for (int[] edge : blueEdges) {
      final int u = edge[0];
      final int v = edge[1];
      graph[u].add(new Pair<>(v, Color.kBlue));
    }

    for (int step = 0; !q.isEmpty(); ++step)
      for (int sz = q.size(); sz > 0; --sz) {
        final int u = q.peek().getKey();
        Color prevColor = q.poll().getValue();
        ans[u] = ans[u] == -1 ? step : ans[u];
        for (int i = 0; i < graph[u].size(); ++i) {
          Pair<Integer, Color> node = graph[u].get(i);
          final int v = node.getKey();
          Color edgeColor = node.getValue();
          if (v == -1 || edgeColor == prevColor)
            continue;
          q.add(new Pair<>(v, edgeColor));
          // Mark (u, v) as used.
          graph[u].set(i, new Pair<>(-1, edgeColor));
        }
      }

    return ans;
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
from enum import Enum


class Color(Enum):
  kInit = 0
  kRed = 1
  kBlue = 2


class Solution:
  def shortestAlternatingPaths(
      self,
      n: int,
      redEdges: list[list[int]],
      blueEdges: list[list[int]],
  ) -> list[int]:
    ans = [-1] * n
    graph = [[] for _ in range(n)]  # graph[u] := [(v, edgeColor)]
    q = collections.deque([(0, Color.kInit)])  # [(u, prevColor)]

    for u, v in redEdges:
      graph[u].append((v, Color.kRed))

    for u, v in blueEdges:
      graph[u].append((v, Color.kBlue))

    step = 0
    while q:
      for _ in range(len(q)):
        u, prevColor = q.popleft()
        if ans[u] == -1:
          ans[u] = step
        for i, (v, edgeColor) in enumerate(graph[u]):
          if v == -1 or edgeColor == prevColor:
            continue
          q.append((v, edgeColor))
          graph[u][i] = (-1, edgeColor)  # Mark (u, v) as used.
      step += 1

    return ans