Maximize the Number of Target Nodes After Connecting Trees I
Source: leetcode
Difficulty: Medium
Tags: tree, depth-first-search, breadth-first-search

There exist two **undirected** trees with `n` and `m` nodes, with **distinct** labels in ranges `[0, n - 1]` and `[0, m - 1]`, respectively.


You are given two 2D integer arrays `edges1` and `edges2` of lengths `n - 1` and `m - 1`, respectively, where `edges1[i] = [ai, bi]` indicates that there is an edge between nodes `ai` and `bi` in the first tree and `edges2[i] = [ui, vi]` indicates that there is an edge between nodes `ui` and `vi` in the second tree. You are also given an integer `k`.


Node `u` is **target** to node `v` if the number of edges on the path from `u` to `v` is less than or equal to `k`. **Note** that a node is *always* **target** to itself.


Return an array of `n` integers `answer`, where `answer[i]` is the **maximum** possible number of nodes **target** to node `i` of the first tree if you have to connect one node from the first tree to another node in the second tree.


**Note** that queries are independent from each other. That is, for every query you will remove the added edge before proceeding to the next query.


 


**Example 1:**



**Input:** edges1 \= \[\[0,1],\[0,2],\[2,3],\[2,4]], edges2 \= \[\[0,1],\[0,2],\[0,3],\[2,7],\[1,4],\[4,5],\[4,6]], k \= 2


**Output:** \[9,7,9,8,8]


**Explanation:**


- For `i = 0`, connect node 0 from the first tree to node 0 from the second tree.
- For `i = 1`, connect node 1 from the first tree to node 0 from the second tree.
- For `i = 2`, connect node 2 from the first tree to node 4 from the second tree.
- For `i = 3`, connect node 3 from the first tree to node 4 from the second tree.
- For `i = 4`, connect node 4 from the first tree to node 4 from the second tree.


![](https://assets.leetcode.com/uploads/2024/09/24/3982-1.png)
**Example 2:**



**Input:** edges1 \= \[\[0,1],\[0,2],\[0,3],\[0,4]], edges2 \= \[\[0,1],\[1,2],\[2,3]], k \= 1


**Output:** \[6,3,3,3,3]


**Explanation:**


For every `i`, connect node `i` of the first tree with any node of the second tree.


![](https://assets.leetcode.com/uploads/2024/09/24/3928-2.png)
 


**Constraints:**


- `2 <= n, m <= 1000`
- `edges1.length == n - 1`
- `edges2.length == m - 1`
- `edges1[i].length == edges2[i].length == 2`
- `edges1[i] = [ai, bi]`
- `0 <= ai, bi < n`
- `edges2[i] = [ui, vi]`
- `0 <= ui, vi < m`
- The input is generated such that `edges1` and `edges2` represent valid trees.
- `0 <= k <= 1000`




C++ solution:
class Solution {
public:
    vector<int> maxTargetNodes(vector<vector<int>>& edges1, vector<vector<int>>& edges2, int k) {
        // Base case: if k is 0, all nodes can only reach themselves.
        if (k == 0) {
            return vector<int>(edges1.size() + 1, 1);
        }

        // Initialize adjacency lists for both trees
        int n1 = edges1.size() + 1, n2 = edges2.size() + 1;
        vector<vector<int>> adjacencyList1(n1), adjacencyList2(n2);

        // Build adjacency list for the first tree
        for (auto& edge : edges1) {
            adjacencyList1[edge[0]].push_back(edge[1]);
            adjacencyList1[edge[1]].push_back(edge[0]);
        }

        // Build adjacency list for the second tree
        for (auto& edge : edges2) {
            adjacencyList2[edge[0]].push_back(edge[1]);
            adjacencyList2[edge[1]].push_back(edge[0]);
        }

        // Function to calculate reachable nodes within a given distance for a tree
        auto calculateReachableNodes_BFS = [](vector<vector<int>>& adjacencyList, int maxDistance) {
            int n = adjacencyList.size();
            vector<int> reachCounts(n, 0);
            
            for (int startNode = 0; startNode < n; ++startNode) {
                vector<bool> visited(n, false);
                queue<int> q;
                q.push(startNode);
                visited[startNode] = true;
                int level = 0;

                while (!q.empty() && level < maxDistance) {
                    int lenQ = q.size();
                    ++level;

                    for (int i = 0; i < lenQ; ++i) {
                        int currentNode = q.front();
                        q.pop();
                        
                        for (int neighbor : adjacencyList[currentNode]) {
                            if (!visited[neighbor]) {
                                visited[neighbor] = true;
                                q.push(neighbor);
                            }
                        }
                    }
                }
                reachCounts[startNode] = count(visited.begin(), visited.end(), true);
            }
            return reachCounts;
        };

        // Calculate reachable nodes for the first tree within distance k
        vector<int> reachableCount1 = calculateReachableNodes_BFS(adjacencyList1, k);

        // Calculate reachable nodes for the second tree within distance k-1
        vector<int> reachableCount2 = calculateReachableNodes_BFS(adjacencyList2, k - 1);

        // Find the maximum reachable nodes in the second tree
        int maxReachableInTree2 = *max_element(reachableCount2.begin(), reachableCount2.end());

        // Add the maximum from the second tree to each node's reachable count in the first tree
        for (int i = 0; i < reachableCount1.size(); ++i) {
            reachableCount1[i] += maxReachableInTree2;
        }

        return reachableCount1;
    }
};

Go solution:
func maxTargetNodes(edges1 [][]int, edges2 [][]int, k int) []int {
	n := len(edges1) + 1
	m := len(edges2) + 1
	graph1 := make([][]int, n)
	for _, e := range edges1 {
		u, v := e[0], e[1]
		graph1[u] = append(graph1[u], v)
		graph1[v] = append(graph1[v], u)
	}
	graph2 := make([][]int, m)
	for _, e := range edges2 {
		u, v := e[0], e[1]
		graph2[u] = append(graph2[u], v)
		graph2[v] = append(graph2[v], u)
	}
	sortedDist := make([][]int, m)
	for start := 0; start < m; start++ {
		dist := make([]int, m)
		for i := 0; i < m; i++ {
			dist[i] = -1
		}
		queue := []int{start}
		dist[start] = 0
		for len(queue) > 0 {
			cur := queue[0]
			queue = queue[1:]
			for _, nxt := range graph2[cur] {
				if dist[nxt] == -1 {
					dist[nxt] = dist[cur] + 1
					queue = append(queue, nxt)
				}
			}
		}
		tmp := make([]int, m)
		for i := 0; i < m; i++ {
			tmp[i] = dist[i]
		}
		for i := 1; i < m; i++ {
			j := i
			for j > 0 && tmp[j] < tmp[j-1] {
				tmp[j], tmp[j-1] = tmp[j-1], tmp[j]
				j--
			}
		}
		sortedDist[start] = tmp
	}
	F := make([]int, k+1)
	for r := 0; r <= k; r++ {
		maxCount := 0
		for i := 0; i < m; i++ {
			lo, hi := 0, m
			for lo < hi {
				mid := (lo + hi) / 2
				if sortedDist[i][mid] <= r {
					lo = mid + 1
				} else {
					hi = mid
				}
			}
			if lo > maxCount {
				maxCount = lo
			}
		}
		F[r] = maxCount
	}
	ans := make([]int, n)
	for start := 0; start < n; start++ {
		dist := make([]int, n)
		for i := 0; i < n; i++ {
			dist[i] = -1
		}
		queue := []int{start}
		dist[start] = 0
		for len(queue) > 0 {
			cur := queue[0]
			queue = queue[1:]
			for _, nxt := range graph1[cur] {
				if dist[nxt] == -1 {
					dist[nxt] = dist[cur] + 1
					queue = append(queue, nxt)
				}
			}
		}
		countT1, bestExtra := 0, 0
		for i := 0; i < n; i++ {
			if dist[i] != -1 && dist[i] <= k {
				countT1++
			}
			if dist[i] != -1 && dist[i]+1 <= k {
				rem := k - (dist[i] + 1)
				if rem >= 0 && F[rem] > bestExtra {
					bestExtra = F[rem]
				}
			}
		}
		ans[start] = countT1 + bestExtra
	}
	return ans
}

Java solution:
class Solution {
    public int[] maxTargetNodes(int[][] edges1, int[][] edges2, int k) {
        if (k == 0) {
            int[] result = new int[edges1.length + 1];
            Arrays.fill(result, 1);
            return result;
        }

        // Create adjacency lists
        List<List<Integer>> adj1 = new ArrayList<>(edges1.length + 1);
        List<List<Integer>> adj2 = new ArrayList<>(edges2.length + 1);

        // Initialize adjacency lists
        for (int i = 0; i <= edges1.length; i++) {
            adj1.add(new ArrayList<>());
        }
        for (int i = 0; i <= edges2.length; i++) {
            adj2.add(new ArrayList<>());
        }

        // Populate adjacency list for first graph
        for (int[] edge : edges1) {
            int u = edge[0], v = edge[1];
            adj1.get(u).add(v);
            adj1.get(v).add(u);
        }

        // Populate adjacency list for second graph
        for (int[] edge : edges2) {
            int u = edge[0], v = edge[1];
            adj2.get(u).add(v);
            adj2.get(v).add(u);
        }

        // Calculate target nodes for first graph
        int[] ans1 = new int[edges1.length + 1];
        for (int i = 0; i < adj1.size(); i++) {
            ans1[i] = bfs(adj1, i, k);
        }

        // Calculate target nodes for second graph
        int[] ans2 = new int[edges2.length + 1];
        for (int i = 0; i < adj2.size(); i++) {
            ans2[i] = bfs(adj2, i, k - 1);
        }

        // Find maximum value in ans2
        int maxi = Arrays.stream(ans2).max().orElse(0);

        // Combine results
        for (int i = 0; i < ans1.length; i++) {
            ans1[i] += maxi;
        }

        return ans1;
    }

    private int bfs(List<List<Integer>> adj, int start, int maxDepth) {
        Queue<Integer> q = new LinkedList<>();
        boolean[] visited = new boolean[adj.size()];
        q.offer(start);
        visited[start] = true;

        int depth = 1;
        while (!q.isEmpty() && depth <= maxDepth) {
            int size = q.size();
            for (int i = 0; i < size; i++) {
                int current = q.poll();
                for (int neighbor : adj.get(current)) {
                    if (!visited[neighbor]) {
                        visited[neighbor] = true;
                        q.offer(neighbor);
                    }
                }
            }
            depth++;
        }

        // Count visited nodes
        int visitedCount = 0;
        for (boolean v : visited) {
            if (v) visitedCount++;
        }
        return visitedCount;
    }
}

JavaScript solution:
/**
 * @param {number[][]} edges1
 * @param {number[][]} edges2
 * @param {number} k
 * @return {number[]}
 */
var maxTargetNodes = function(edges1, edges2, k) {
    const n = edges1.length + 1;
    const m = edges2.length + 1;

    const adj1 = Array.from({ length: n }, () => []);
    const adj2 = Array.from({ length: m }, () => []);

    for (const [u, v] of edges1) {
      adj1[u].push(v);
      adj1[v].push(u);
    }
    for (const [u, v] of edges2) {
      adj2[u].push(v);
      adj2[v].push(u);
    } 
    const good1 = new Array(n).fill(0);
    const good2 = new Array(m).fill(0);

    const dfs = (node, parent, distance, root, k, good, adj) => {
      if (distance >= k) return;
      good[root]++;
      for (const neighbor of adj[node]) {
        if (neighbor !== parent) {
          dfs(neighbor, node, distance + 1, root, k, good, adj);
        }
      }
    };
    for (let i = 0; i < n; i++) {
      dfs(i, -1, 0, i, k + 1, good1, adj1);
    }

    for (let i = 0; i < m; i++) {
      dfs(i, -1, 0, i, k, good2, adj2);
    }

    const mx = Math.max(...good2);

    return good1.map(value => value + mx);

};

Python solution:
class Solution:
    def maxTargetNodes(self, edges1: List[List[int]], edges2: List[List[int]], k: int) -> List[int]:
        
        def fn(edges, k): 
            n = len(edges)+1
            tree = [[] for _ in range(n)]
            for u, v in edges: 
                tree[u].append(v)
                tree[v].append(u)
            ans = [0]*n
            for x in range(n): 
                stack = [(x, -1, 0)]
                while stack: 
                    u, p, d = stack.pop()
                    if d <= k: ans[x] += 1
                    for v in tree[u]: 
                        if v != p: 
                            stack.append((v, u, d+1))
            return ans 
        
        most = max(fn(edges2, k-1))
        return [x+most for x in fn(edges1, k)]

Ruby solution:
def max_target_nodes(edges1, edges2, k)
  n = edges1.size + 1
  m = edges2.size + 1
  adj1 = Array.new(n){ [] }
  edges1.each { |a,b| adj1[a] << b; adj1[b] << a }
  adj2 = Array.new(m){ [] }
  edges2.each { |u,v| adj2[u] << v; adj2[v] << u }
  if k == 0
    return [1] * n
  end
  bfs_count = lambda do |start, limit, adj|
    size = adj.size
    visited = Array.new(size, false)
    dist = Array.new(size, 0)
    queue = [start]
    visited[start] = true
    head = 0
    cnt = 0
    while head < queue.size
      cur = queue[head]
      head += 1
      cnt += 1
      if dist[cur] < limit
        for nb in adj[cur]
          unless visited[nb]
            visited[nb] = true
            dist[nb] = dist[cur] + 1
            queue << nb
          end
        end
      end
    end
    cnt
  end
  extra = 0
  (0...m).each do |v|
    cnt = bfs_count.call(v, k - 1, adj2)
    extra = cnt if cnt > extra
  end
  ans = Array.new(n, 0)
  (0...n).each do |i|
    ans[i] = bfs_count.call(i, k, adj1) + extra
  end
  ans
end