Skip to content

3373. Maximize the Number of Target Nodes After Connecting Trees II 👍

  • Time: $O(n + m)$
  • Space: $O(n + m)$
 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
47
48
49
50
class Solution {
 public:
  vector<int> maxTargetNodes(vector<vector<int>>& edges1,
                             vector<vector<int>>& edges2) {
    const int n = edges1.size() + 1;
    const int m = edges2.size() + 1;
    vector<int> ans;
    const vector<vector<int>> graph1 = buildGraph(edges1);
    const vector<vector<int>> graph2 = buildGraph(edges2);
    vector<bool> parity1(n);
    vector<bool> parity2(m);  // placeholder (parity2 is not used)
    const int even1 = dfs(graph1, 0, -1, parity1, /*isEven=*/true);
    const int even2 = dfs(graph2, 0, -1, parity2, /*isEven=*/true);
    const int odd1 = n - even1;
    const int odd2 = m - even2;

    for (int i = 0; i < n; ++i) {
      const int tree1 = parity1[i] ? even1 : odd1;
      // Can connect the current node in tree1 to either an even node or an odd
      // node in tree2.
      const int tree2 = max(even2, odd2);
      ans.push_back(tree1 + tree2);
    }

    return ans;
  }

 private:
  // Returns the number of nodes that can be reached from u with even steps.
  int dfs(const vector<vector<int>>& graph, int u, int prev,
          vector<bool>& parity, bool isEven) {
    int res = isEven ? 1 : 0;
    parity[u] = isEven;
    for (const int v : graph[u])
      if (v != prev)
        res += dfs(graph, v, u, parity, !isEven);
    return res;
  }

  vector<vector<int>> buildGraph(const vector<vector<int>>& edges) {
    vector<vector<int>> graph(edges.size() + 1);
    for (const vector<int>& edge : edges) {
      const int u = edge[0];
      const int v = edge[1];
      graph[u].push_back(v);
      graph[v].push_back(u);
    }
    return graph;
  }
};

 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
47
48
class Solution {
  public int[] maxTargetNodes(int[][] edges1, int[][] edges2) {
    final int n = edges1.length + 1;
    final int m = edges2.length + 1;
    List<Integer>[] graph1 = buildGraph(edges1, n);
    List<Integer>[] graph2 = buildGraph(edges2, m);
    boolean[] parity1 = new boolean[n];
    boolean[] parity2 = new boolean[m]; // Placeholder (not used)
    final int even1 = dfs(graph1, 0, -1, parity1, true);
    final int even2 = dfs(graph2, 0, -1, parity2, true);
    final int odd1 = n - even1;
    final int odd2 = m - even2;
    int[] ans = new int[n];

    for (int i = 0; i < n; i++) {
      final int tree1 = parity1[i] ? even1 : odd1;
      // Can connect the current node in tree1 to either an even or an odd node
      // in tree2.
      final int tree2 = Math.max(even2, odd2);
      ans[i] = tree1 + tree2;
    }

    return ans;
  }

  // Returns the number of nodes that can be reached from u with even steps.
  private int dfs(List<Integer>[] graph, int u, int prev, boolean[] parity, boolean isEven) {
    int res = isEven ? 1 : 0;
    parity[u] = isEven;
    for (final int v : graph[u])
      if (v != prev)
        res += dfs(graph, v, u, parity, !isEven);
    return res;
  }

  private List<Integer>[] buildGraph(int[][] edges, int n) {
    List<Integer>[] graph = new ArrayList[n];
    for (int i = 0; i < n; i++)
      graph[i] = new ArrayList<>();
    for (int[] edge : edges) {
      final int u = edge[0];
      final int v = edge[1];
      graph[u].add(v);
      graph[v].add(u);
    }
    return graph;
  }
}

 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
class Solution:
  def maxTargetNodes(
      self,
      edges1: list[list[int]],
      edges2: list[list[int]]
  ) -> list[int]:
    n = len(edges1) + 1
    m = len(edges2) + 1
    graph1 = self._buildGraph(edges1)
    graph2 = self._buildGraph(edges2)
    parity1 = [False] * n
    parity2 = [False] * m  # placeholder (parity2 is not used)
    even1 = self._dfs(graph1, 0, -1, parity1, True)
    even2 = self._dfs(graph2, 0, -1, parity2, True)
    odd1 = n - even1
    odd2 = m - even2

    # Can connect the current node in tree1 to either an even node or an odd
    # node in tree2.
    return [(even1 if parity1[i] else odd1) + max(even2, odd2)
            for i in range(n)]

  def _dfs(
      self,
      graph: list[list[int]],
      u: int,
      prev: int,
      parity: list[bool],
      isEven: bool
  ) -> int:
    """
    Returns the number of nodes that can be reached from u with even steps.
    """
    res = 1 if isEven else 0
    parity[u] = isEven
    for v in graph[u]:
      if v != prev:
        res += self._dfs(graph, v, u, parity, not isEven)
    return res

  def _buildGraph(self, edges: list[list[int]]) -> list[list[int]]:
    graph = [[] for _ in range(len(edges) + 1)]
    for u, v in edges:
      graph[u].append(v)
      graph[v].append(u)
    return graph