Medium

题目描述

存在两个无向树，分别有 n 和 m 个节点，节点标签分别在范围 [0, n - 1] 和 [0, m - 1] 内。

给你两个长度分别为 n - 1 和 m - 1 的二维整数数组 edges1 和 edges2，其中 edges1[i] = [ai, bi] 表示第一棵树中节点 ai 和 bi 之间有一条边，edges2[i] = [ui, vi] 表示第二棵树中节点 ui 和 vi 之间有一条边。同时给你一个整数 k。

如果从节点 u 到节点 v 的路径上的边数小于等于 k，则称节点 u 是节点 v 的目标节点。注意，节点总是自己的目标节点。

返回一个长度为 n 的整数数组 answer，其中 answer[i] 是在将第一棵树的一个节点连接到第二棵树的一个节点后，第一棵树中节点 i 的目标节点的最大可能数量。

注意，查询是相互独立的。也就是说，对于每个查询，你都会在进行下一个查询之前移除添加的边。

示例 1：

输入：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
输出：[9,7,9,8,8]

示例 2：

输入：edges1 = [[0,1],[0,2],[0,3],[0,4]], edges2 = [[0,1],[1,2],[2,3]], k = 1
输出：[6,3,3,3,3]

提示：

  • 2 <= n, m <= 1000
  • edges1.length == n - 1
  • edges2.length == m - 1
  • 0 <= k <= 1000

解题思路

解题思路

这道题的核心是理解如何通过在两棵树之间添加一条边来最大化目标节点数量。

关键观察：

  1. 对于第一棵树中的每个节点 i，我们需要计算距离它不超过 k 的所有节点数量
  2. 当我们在两棵树之间添加一条边时，第一棵树中的节点还可以通过这条边到达第二棵树中的节点
  3. 如果第一棵树的节点 u 连接到第二棵树的节点 v，那么从第一棵树节点 i 到第二棵树节点 j 的距离为：dist(i,u) + 1 + dist(v,j)

算法步骤：

  1. 对于第一棵树，使用 BFS/DFS 计算每个节点到其他所有节点的距离，统计距离不超过 k 的节点数
  2. 对于第二棵树，计算每个节点到其他所有节点距离不超过 k-1 的节点数（因为要额外花费 1 的距离通过连接边）
  3. 找到第二棵树中能提供最大额外目标节点数的节点
  4. 对于第一棵树的每个节点 i，答案是它在原树中的目标节点数加上通过最优连接能获得的额外目标节点数

时间复杂度分析：

  • 计算所有节点对距离：O(n² + m²)
  • 由于 n,m ≤ 1000，这个复杂度是可接受的

代码实现

class Solution {
public:
    vector<int> maxTargetNodes(vector<vector<int>>& edges1, vector<vector<int>>& edges2, int k) {
        int n = edges1.size() + 1;
        int m = edges2.size() + 1;
        
        // Build adjacency lists
        vector<vector<int>> adj1(n), adj2(m);
        for (auto& edge : edges1) {
            adj1[edge[0]].push_back(edge[1]);
            adj1[edge[1]].push_back(edge[0]);
        }
        for (auto& edge : edges2) {
            adj2[edge[0]].push_back(edge[1]);
            adj2[edge[1]].push_back(edge[0]);
        }
        
        // Calculate distances and count target nodes for tree1
        vector<int> count1(n);
        for (int i = 0; i < n; i++) {
            count1[i] = countNodes(adj1, i, k);
        }
        
        // Calculate distances and count target nodes for tree2
        int maxCount2 = 0;
        if (k > 0) {
            for (int i = 0; i < m; i++) {
                maxCount2 = max(maxCount2, countNodes(adj2, i, k - 1));
            }
        }
        
        // Build result
        vector<int> result(n);
        for (int i = 0; i < n; i++) {
            result[i] = count1[i] + maxCount2;
        }
        
        return result;
    }
    
private:
    int countNodes(vector<vector<int>>& adj, int start, int maxDist) {
        if (maxDist < 0) return 0;
        
        int n = adj.size();
        vector<int> dist(n, -1);
        queue<int> q;
        
        dist[start] = 0;
        q.push(start);
        int count = 1;
        
        while (!q.empty()) {
            int node = q.front();
            q.pop();
            
            if (dist[node] >= maxDist) continue;
            
            for (int neighbor : adj[node]) {
                if (dist[neighbor] == -1) {
                    dist[neighbor] = dist[node] + 1;
                    q.push(neighbor);
                    count++;
                }
            }
        }
        
        return count;
    }
};

class Solution:
    def maxTargetNodes(self, edges1: List[List[int]], edges2: List[List[int]], k: int) -> List[int]:
        from collections import defaultdict, deque
        
        n = len(edges1) + 1
        m = len(edges2) + 1
        
        # Build adjacency lists
        adj1 = defaultdict(list)
        adj2 = defaultdict(list)
        
        for a, b in edges1:
            adj1[a].append(b)
            adj1[b].append(a)
            
        for a, b in edges2:
            adj2[a].append(b)
            adj2[b].append(a)
        
        def count_nodes(adj, start, max_dist):
            if max_dist < 0:
                return 0
                
            dist = [-1] * len(adj)
            queue = deque([start])
            dist[start] = 0
            count = 1
            
            while queue:
                node = queue.popleft()
                
                if dist[node] >= max_dist:
                    continue
                
                for neighbor in adj[node]:
                    if dist[neighbor] == -1:
                        dist[neighbor] = dist[node] + 1
                        queue.append(neighbor)
                        count += 1
            
            return count
        
        # Calculate target nodes for each node in tree1
        count1 = [count_nodes(adj1, i, k) for i in range(n)]
        
        # Find maximum target nodes from tree2 (with k-1 distance)
        max_count2 = 0
        if k > 0:
            for i in range(m):
                max_count2 = max(max_count2, count_nodes(adj2, i, k - 1))
        
        # Build result
        return [count1[i] + max_count2 for i in range(n)]

public class Solution {
    public int[] MaxTargetNodes(int[][] edges1, int[][] edges2, int k) {
        int n = edges1.Length + 1;
        int m = edges2.Length + 1;
        
        // Build adjacency lists
        List<int>[] adj1 = new List<int>[n];
        List<int>[] adj2 = new List<int>[m];
        
        for (int i = 0; i < n; i++) adj1[i] = new List<int>();
        for (int i = 0; i < m; i++) adj2[i] = new List<int>();
        
        foreach (var edge in edges1) {
            adj1[edge[0]].Add(edge[1]);
            adj1[edge[1]].Add(edge[0]);
        }
        
        foreach (var edge in edges2) {
            adj2[edge[0]].Add(edge[1]);
            adj2[edge[1]].Add(edge[0]);
        }
        
        // Calculate target nodes for each node in tree1
        int[] count1 = new int[n];
        for (int i = 0; i < n; i++) {
            count1[i] = CountNodes(adj1, i, k);
        }
        
        // Find maximum target nodes from tree2
        int maxCount2 = 0;
        if (k > 0) {
            for (int i = 0; i < m; i++) {
                maxCount2 = Math.Max(maxCount2, CountNodes(adj2, i, k - 1));
            }
        }
        
        // Build result
        int[] result = new int[n];
        for (int i = 0; i < n; i++) {
            result[i] = count1[i] + maxCount2;
        }
        
        return result;
    }
    
    private int CountNodes(List<int>[] adj, int start, int maxDist) {
        if (maxDist < 0) return 0;
        
        int n = adj.Length;
        int[] dist = new int[n];
        Array.Fill(dist, -1);
        
        Queue<int> queue = new Queue<int>();
        dist[start] = 0;
        queue.Enqueue(start);
        int count = 1;
        
        while (queue.Count > 0) {
            int node = queue.Dequeue();
            
            if (dist[node] >= maxDist) continue;
            
            foreach (int neighbor in adj[node]) {
                if (dist[neighbor] == -1) {
                    dist[neighbor] = dist[node] + 1;
                    queue.Enqueue(neighbor);
                    count++;
                }
            }
        }
        
        return count;
    }
}

var maxTargetNodes = function(edges1, edges2, k) {
    function buildGraph(edges) {
        const graph = {};
        for (const [u, v] of edges) {
            if (!graph[u]) graph[u] = [];
            if (!graph[v]) graph[v] = [];
            graph[u].push(v);
            graph[v].push(u);
        }
        return graph;
    }
    
    function countTargets(graph, start, maxDist) {
        const visited = new Set();
        const queue = [[start, 0]];
        visited.add(start);
        let count = 0;
        
        while (queue.length > 0) {
            const [node, dist] = queue.shift();
            if (dist <= maxDist) {
                count++;
                if (graph[node]) {
                    for (const neighbor of graph[node]) {
                        if (!visited.has(neighbor)) {
                            visited.add(neighbor);
                            queue.push([neighbor, dist + 1]);
                        }
                    }
                }
            }
        }
        return count;
    }
    
    const graph1 = buildGraph(edges1);
    const graph2 = buildGraph(edges2);
    const n = edges1.length + 1;
    const m = edges2.length + 1;
    
    // For each node in tree2, calculate how many nodes are reachable within k-1 distance
    let maxTree2Contribution = 0;
    for (let j = 0; j < m; j++) {
        const contribution = countTargets(graph2, j, k - 1);
        maxTree2Contribution = Math.max(maxTree2Contribution, contribution);
    }
    
    const result = [];
    
    // For each node in tree1
    for (let i = 0; i < n; i++) {
        const tree1Count = countTargets(graph1, i, k);
        result.push(tree1Count + maxTree2Contribution);
    }
    
    return result;
};
在 LeetCode 上运行： 力扣中文版 LeetCode 国际版

复杂度分析

复杂度类型复杂度说明
时间复杂度O(n² + m²)对每个节点进行BFS遍历整棵树
空间复杂度O(n + m)邻接表存储和BFS队列空间

相关题目