Maximize the Number of Target Nodes After Connecting Trees I

Med

#2996Time: O(n^2 + m^2 + n^2 * m). `O(n^2)` for `dist1`, `O(m^2)` for `dist2` and `count2`. The main calculation involves a loop for `i` (n), a nested loop for `p` (n), and an inner loop for `q` (m), resulting in `O(n^2 * m)` for the main logic.Space: O(n^2 + m^2) to store the all-pairs shortest path matrices for both trees and the precomputed counts for the second tree.1 company

Algorithms

Depth-First Search Breadth-First Search

Data structures

Tree

Companies

jio

Prompt

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] = [a_i, b_i] indicates that there is an edge between nodes a_i and b_i in the first tree and edges2[i] = [u_i, v_i] indicates that there is an edge between nodes u_i and v_i 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.

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.

Constraints:

2 <= n, m <= 1000
edges1.length == n - 1
edges2.length == m - 1
edges1[i].length == edges2[i].length == 2
edges1[i] = [a_i, b_i]
0 <= a_i, b_i < n
edges2[i] = [u_i, v_i]
0 <= u_i, v_i < m
The input is generated such that edges1 and edges2 represent valid trees.
0 <= k <= 1000

Approaches

2 approaches with complexity analysis and trade-offs.

This approach involves iterating through all possible connection points to find the maximum number of target nodes. For each node i in the first tree, we consider connecting every node p from the first tree to every node q from the second tree. We calculate the total number of target nodes for this specific connection and keep track of the maximum value found. To avoid re-computation, we can precompute all-pairs shortest paths and reachable node counts for both trees.

Algorithm

Build adjacency lists for both trees. - Precompute all-pairs shortest paths dist1 for tree 1 using n BFS runs. - Precompute all-pairs shortest paths dist2 for tree 2 using m BFS runs. - Precompute a table count2[q][r] storing the number of nodes in tree 2 within distance r from node q. - For each node i in tree 1: - Calculate count_in_tree1, the number of nodes in tree 1 within distance k of i. - Initialize max_from_tree2 = 0. - Iterate through all possible connection nodes p in tree 1. - For each p, determine the required radius in tree 2: radius = k - 1 - dist1[i][p]. - If radius >= 0, find the best connection node q in tree 2 for this p. This involves finding max_{q \in V2} count2[q][radius]. - Update max_from_tree2 with the value found. - answer[i] = count_in_tree1 + max_from_tree2.

Walkthrough

The core of this method is a brute-force search over the connection choices. For each query i, we want to find the best pair of nodes (p, q) to connect, where p is from the first tree and q is from the second. The total number of target nodes for i with this connection is the sum of targets in tree 1 and tree 2. The number of targets in tree 1 is fixed for a given i. The number of targets in tree 2 depends on the connection (p, q). The path from i to a node v2 in tree 2 goes through p and q, so its length is dist1(i, p) + 1 + dist2(q, v2). We need this to be <= k. This means dist2(q, v2) <= k - 1 - dist1(i, p). We iterate through all p and q to maximize the count of such v2 nodes. Precomputing distances and counts helps, but the triple loop over i, p, and q remains the bottleneck.

1java class Solution {     public int[] maximizeTarget(int n, int[][] edges1, int m, int[][] edges2, int k) {         List<List<Integer>> adj1 = buildAdj(n, edges1);         List<List<Integer>> adj2 = buildAdj(m, edges2);          int[][] dist1 = computeAllPairsDist(n, adj1);         int[][] dist2 = computeAllPairsDist(m, adj2);          // Precompute counts for tree2 for each radius         int[][] count2 = new int[m][m];         for (int i = 0; i < m; i++) {             for (int j = 0; j < m; j++) {                 if (dist2[i][j] < m) {                     count2[i][dist2[i][j]]++;                 }             }             for (int r = 1; r < m; r++) {                 count2[i][r] += count2[i][r - 1];             }         }          int[] answer = new int[n];         for (int i = 0; i < n; i++) {             int countInTree1 = 0;             for (int v1 = 0; v1 < n; v1++) {                 if (dist1[i][v1] <= k) {                     countInTree1++;                 }             }              int maxFromTree2 = 0;             for (int p = 0; p < n; p++) {                 int radius = k - 1 - dist1[i][p];                 if (radius >= 0) {                     int currentMaxForP = 0;                     for (int q = 0; q < m; q++) {                         currentMaxForP = Math.max(currentMaxForP, count2[q][Math.min(radius, m - 1)]);                     }                     maxFromTree2 = Math.max(maxFromTree2, currentMaxForP);                 }             }             answer[i] = countInTree1 + maxFromTree2;         }         return answer;     }      private List<List<Integer>> buildAdj(int size, int[][] edges) {         List<List<Integer>> adj = new ArrayList<>();         for (int i = 0; i < size; i++) adj.add(new ArrayList<>());         for (int[] edge : edges) {             adj.get(edge[0]).add(edge[1]);             adj.get(edge[1]).add(edge[0]);         }         return adj;     }      private int[][] computeAllPairsDist(int size, List<List<Integer>> adj) {         int[][] dist = new int[size][size];         for (int i = 0; i < size; i++) {             Arrays.fill(dist[i], -1);             Queue<Integer> q = new LinkedList<>();             q.offer(i);             dist[i][i] = 0;             int d = 1;             while (!q.isEmpty()) {                 int levelSize = q.size();                 for (int j = 0; j < levelSize; j++) {                     int u = q.poll();                     for (int v : adj.get(u)) {                         if (dist[i][v] == -1) {                             dist[i][v] = d;                             q.offer(v);                         }                     }                 }                 d++;             }         }         return dist;     } }

Complexity

Time

O(n^2 + m^2 + n^2 * m). `O(n^2)` for `dist1`, `O(m^2)` for `dist2` and `count2`. The main calculation involves a loop for `i` (n), a nested loop for `p` (n), and an inner loop for `q` (m), resulting in `O(n^2 * m)` for the main logic.

Space

O(n^2 + m^2) to store the all-pairs shortest path matrices for both trees and the precomputed counts for the second tree.

Trade-offs

Pros

It is a direct and straightforward implementation based on the problem's definition.
It correctly solves the problem for small inputs.

Cons

The time complexity of O(n^2 * m) is too high for the given constraints, leading to a 'Time Limit Exceeded' verdict on most platforms.

This approach improves upon the brute-force method by making a key observation. For a fixed node i in the first tree, to maximize the number of target nodes from the second tree, we must maximize the allowed radius for searching in the second tree. The radius is k - 1 - dist1(i, p), where p is the connection point in the first tree. This expression is maximized when dist1(i, p) is minimized. The minimum possible distance dist1(i, p) is 0, which occurs when we choose p = i. Therefore, for each query i, the optimal connection point in the first tree is node i itself. This simplifies the problem significantly, removing one layer of iteration from the main loop.

Algorithm

The core insight is that for any query node i, the optimal connection point p in the first tree is always i itself. This maximizes the available distance budget for the second tree. - The problem simplifies to answer[i] = count1(i) + max_count_from_tree2. - count1(i) is the number of nodes in tree 1 within distance k of i. - max_count_from_tree2 is the maximum number of nodes in tree 2 reachable from any single node q within a radius of k-1. This value is constant for all i. - Precomputation for Tree 1: - Build adjacency list and compute all-pairs shortest paths dist1. - For each i, compute count1[i]. - Precomputation for Tree 2: - Build adjacency list and compute all-pairs shortest paths dist2. - Calculate max_count_from_tree2 by checking each node q as a potential center and finding the one that covers the most nodes within radius k-1. - Final Calculation: - For each i, compute answer[i] = count1[i] + max_count_from_tree2.

Walkthrough

With the observation that the best connection point p in tree 1 is always the query node i, the calculation for answer[i] becomes much simpler. The total number of targets is the sum of targets in tree 1 and tree 2. 1. Targets in Tree 1: The number of nodes v1 where dist1(i, v1) <= k. We can precompute this for all i and store it in an array count1. 2. Targets in Tree 2: When connecting (i, q), the path to a node v2 in tree 2 has length dist1(i, i) + 1 + dist2(q, v2) = 1 + dist2(q, v2). For v2 to be a target, we need 1 + dist2(q, v2) <= k, or dist2(q, v2) <= k-1. To maximize this count, we must choose the node q in tree 2 that has the most neighbors within radius k-1. This maximum count is a single value that we can precompute. The final answer for i is then count1[i] plus this precomputed maximum from tree 2.

1java class Solution {     public int[] maximizeTarget(int n, int[][] edges1, int m, int[][] edges2, int k) {         // Precomputation for Tree 1         List<List<Integer>> adj1 = buildAdj(n, edges1);         int[][] dist1 = computeAllPairsDist(n, adj1);         int[] count1 = new int[n];         for (int i = 0; i < n; i++) {             for (int v = 0; v < n; v++) {                 if (dist1[i][v] <= k) {                     count1[i]++;                 }             }         }          // Precomputation for Tree 2         List<List<Integer>> adj2 = buildAdj(m, edges2);         int[][] dist2 = computeAllPairsDist(m, adj2);                  int maxCountFromTree2 = 0;         if (k - 1 >= 0) {             for (int q = 0; q < m; q++) {                 int currentCount = 0;                 for (int v = 0; v < m; v++) {                     if (dist2[q][v] <= k - 1) {                         currentCount++;                     }                 }                 maxCountFromTree2 = Math.max(maxCountFromTree2, currentCount);             }         }          // Final Calculation         int[] answer = new int[n];         for (int i = 0; i < n; i++) {             answer[i] = count1[i] + maxCountFromTree2;         }                  return answer;     }      private List<List<Integer>> buildAdj(int size, int[][] edges) {         List<List<Integer>> adj = new ArrayList<>();         for (int i = 0; i < size; i++) adj.add(new ArrayList<>());         for (int[] edge : edges) {             adj.get(edge[0]).add(edge[1]);             adj.get(edge[1]).add(edge[0]);         }         return adj;     }      private int[][] computeAllPairsDist(int size, List<List<Integer>> adj) {         int[][] dist = new int[size][size];         for (int i = 0; i < size; i++) {             Arrays.fill(dist[i], -1);             Queue<Integer> q = new LinkedList<>();             q.offer(i);             dist[i][i] = 0;             int d = 1;             while (!q.isEmpty()) {                 int levelSize = q.size();                 for (int j = 0; j < levelSize; j++) {                     int u = q.poll();                     for (int v : adj.get(u)) {                         if (dist[i][v] == -1) {                             dist[i][v] = d;                             q.offer(v);                         }                     }                 }                 d++;             }         }         return dist;     } }

Complexity

Time

O(n^2 + m^2). `O(n^2)` to compute `dist1` and `count1`. `O(m^2)` to compute `dist2` and find the maximum number of reachable nodes in tree 2. The final calculation is a simple `O(n)` loop. The overall complexity is dominated by the APSP precomputation.

Space

O(n^2 + m^2) to store the all-pairs shortest path matrices for both trees.

Trade-offs

Pros

Highly efficient with a time complexity that meets the problem constraints.
The logic is simplified due to the key insight about the optimal connection point.

Cons

Requires O(n^2 + m^2) space for the distance matrices, which might be large for huge values of n and m (though acceptable for the given constraints).

Solutions

Solution

1public class Solution {2    public int[] MaxTargetNodes(int[][] edges1, int[][] edges2, int k) {3        var g2 = Build(edges2);4        int m = edges2.Length + 1;5        int t = 0;6        for (int i = 0; i < m; i++) {7            t = Math.Max(t, Dfs(g2, i, -1, k - 1));8        }9        var g1 = Build(edges1);10        int n = edges1.Length + 1;11        var ans = new int[n];12        Array.Fill(ans, t);13        for (int i = 0; i < n; i++) {14            ans[i] += Dfs(g1, i, -1, k);15        }16        return ans;17    }18    private List < int > [] Build(int[][] edges) {19        int n = edges.Length + 1;20        var g = new List < int > [n];21        for (int i = 0; i < n; i++) {22            g[i] = new List < int > ();23        }24        foreach(var e in edges) {25            int a = e[0], b = e[1];26            g[a].Add(b);27            g[b].Add(a);28        }29        return g;30    }31    private int Dfs(List < int > [] g, int a, int fa, int d) {32        if (d < 0) {33            return 0;34        }35        int cnt = 1;36        foreach(var b in g[a]) {37            if (b != fa) {38                cnt += Dfs(g, b, a, d - 1);39            }40        }41        return cnt;42    }43}

1class Solution {2public3  int[] maxTargetNodes(int[][] edges1, int[][] edges2, int k) {4    var g2 = build(edges2);5    int m = edges2.length + 1;6    int t = 0;7    for (int i = 0; i < m; ++i) {8      t = Math.max(t, dfs(g2, i, -1, k - 1));9    }10    var g1 = build(edges1);11    int n = edges1.length + 1;12    int[] ans = new int[n];13    Arrays.fill(ans, t);14    for (int i = 0; i < n; ++i) {15      ans[i] += dfs(g1, i, -1, k);16    }17    return ans;18  }19private20  List<Integer>[] build(int[][] edges) {21    int n = edges.length + 1;22    List<Integer>[] g = new List[n];23    Arrays.setAll(g, i->new ArrayList<>());24    for (var e : edges) {25      int a = e[0], b = e[1];26      g[a].add(b);27      g[b].add(a);28    }29    return g;30  } private int dfs(List<Integer>[] g, int a, int fa, int d) {31    if (d < 0) {32      return 0;33    }34    int cnt = 1;35    for (int b : g[a]) {36      if (b != fa) {37        cnt += dfs(g, b, a, d - 1);38      }39    }40    return cnt;41  }42}

1class Solution {2public:3  vector<int> maxTargetNodes(vector<vector<int>> &edges1,4                             vector<vector<int>> &edges2, int k) {5    auto g2 = build(edges2);6    int m = edges2.size() + 1;7    int t = 0;8    for (int i = 0; i < m; ++i) {9      t = max(t, dfs(g2, i, -1, k - 1));10    }11    auto g1 = build(edges1);12    int n = edges1.size() + 1;13    vector<int> ans(n, t);14    for (int i = 0; i < n; ++i) {15      ans[i] += dfs(g1, i, -1, k);16    }17    return ans;18  }19 20private:21  vector<vector<int>> build(const vector<vector<int>> &edges) {22    int n = edges.size() + 1;23    vector<vector<int>> g(n);24    for (const auto &e : edges) {25      int a = e[0], b = e[1];26      g[a].push_back(b);27      g[b].push_back(a);28    }29    return g;30  }31  int dfs(const vector<vector<int>> &g, int a, int fa, int d) {32    if (d < 0) {33      return 0;34    }35    int cnt = 1;36    for (int b : g[a]) {37      if (b != fa) {38        cnt += dfs(g, b, a, d - 1);39      }40    }41    return cnt;42  }43};

1class Solution:2    def maxTargetNodes(self, edges1: List[List[int]], edges2: List[List[int]], k: int) -> List[int]: def build(edges: List[List[int]]) -> List[List[int]]: n = len(edges) + 1 g = [[] for _ in range(n)] for a, b in edges: g[a]. append(b) g[b]. append(a) return g def dfs(g: List[List[int]], a: int, fa: int, d: int) -> int: if d < 0: return 0 cnt = 1 for b in g[a]: if b != fa: cnt += dfs(g, b, a, d - 1) return cnt g2 = build(edges2) m = len(edges2) + 1 t = max(dfs(g2, i, - 1, k - 1) for i in range(m)) g1 = build(edges1) n = len(edges1) + 1 return [dfs(g1, i, - 1, k) + t for i in range(n)]

Video walkthrough

Newsletter

One sharp idea, every week

System design and interview prep — short enough to finish.

No spam. Unsubscribe anytime.

Practice

Same difficulty — related problems to reinforce the pattern.

← Previous

Previous problem

Next problem