[LeetCode] 每日一题 3372. 连接两棵树后最大目标节点数目 I（DFS + 思考题）

题目描述

有两棵 无向 树，分别有 n 和 m 个树节点。两棵树中的节点编号分别为[0, n - 1] 和 [0, m - 1] 中的整数。

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

如果节点 u 和节点 v 之间路径的边数小于等于 k ，那么我们称节点 u 是节点 v 的 目标节点 。注意 ，一个节点一定是它自己的 目标节点 。

请你返回一个长度为 n 的整数数组 answer ，answer[i] 表示将第一棵树中的一个节点与第二棵树中的一个节点连接一条边后，第一棵树中节点 i 的 目标节点 数目的 最大值 。

注意 ，每个查询相互独立。意味着进行下一次查询之前，你需要先把刚添加的边给删掉。

题目链接

https://leetcode.cn/problems/maximize-the-number-of-target-nodes-after-connecting-trees-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]

解释：

对于 i = 0 ，连接第一棵树中的节点 0 和第二棵树中的节点 0 。
对于 i = 1 ，连接第一棵树中的节点 1 和第二棵树中的节点 0 。
对于 i = 2 ，连接第一棵树中的节点 2 和第二棵树中的节点 4 。
对于 i = 3 ，连接第一棵树中的节点 3 和第二棵树中的节点 4 。
对于 i = 4 ，连接第一棵树中的节点 4 和第二棵树中的节点 4 。

示例 2

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

输出：[6,3,3,3,3]

解释：

对于每个 i ，连接第一棵树中的节点 i 和第二棵树中的任意一个节点。

提示

• 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

• 输入保证 edges1 和 edges2 都表示合法的树。

• 0 <= k <= 1000

解题思路

这题第一眼看上去让我们对两个树进行连接，然后求目标节点个数的最大值，但仔细一看，操作之间是独立的，也就是说我们不需要真的模拟连接两棵树的过程，而只要找出连接哪一个节点能让目标节点数最多

所谓“目标节点”，是指从某个节点出发，能在不超过 k 条边内到达的所有节点（包含自身）。本质上就是一个限定深度的 DFS 问题

第一步：计算第一棵树中，每个节点自身能到达的目标节点数

很自然想到用 DFS 去遍历所有可能路径。我们从每个节点出发向外深搜，每走一层，k 就减一，直到 k 为负或没有邻居为止。这过程中有个小坑，由于是无向树，如果不做处理可能会在节点之间反复来回跳转。我们可以记录上一个访问的节点，这样在 DFS 时就可以跳过它，避免死循环。

第二步：处理第二棵树的连接点选择

本题的 tricky 点在于连接操作看似复杂，其实可以简化。我们不是要模拟把第一棵树的每个节点都接到第二棵树的所有节点上，而是只需要找到第二棵树中，在 k - 1 步之内最多能到达多少节点的一个最优点。因为每次连接只能走一条边过去，剩下的步数就是 k - 1，因此只要在第二棵树中找到一个节点能在 k - 1 步内尽可能多地遍历，就能让每个第一棵树的节点都尽可能受益

所以只要在第二棵树中每个节点上跑一遍 DFS，找出最大的目标节点数就行了

代码实现

class Solution {
    public int[] maxTargetNodes(int[][] edges1, int[][] edges2, int k) {
        int n = edges1.length + 1, m = edges2.length + 1;
        List<List<Integer>> path1 = getPath(edges1);
        int max = 0;
        if (k > 0) {
            List<List<Integer>> path2 = getPath(edges2);
            for (int i = 0; i < m; i++) {
                max = Math.max(max, dfs(path2, i, -1, k - 1));
            }
        }
        int[] ans = new int[n];
        for (int i = 0; i < n; i++) {
            ans[i] = dfs(path1, i, -1, k) + max;
        }
        return ans;
    }

    private List<List<Integer>> getPath(int[][] edges) {
        int n = edges.length + 1;
        List<List<Integer>> paths = new ArrayList<>();
        for (int i = 0; i < n; i++) {
            paths.add(new ArrayList<>());
        }
        for (int[] edge : edges) {
            paths.get(edge[0]).add(edge[1]);
            paths.get(edge[1]).add(edge[0]);
        }
        return paths;
    }

    private int dfs(List<List<Integer>> path, int index, int pre, int k) {
        if (k < 0) {
            return 0;
        }
        int count = 1;
        for (Integer next : path.get(index)) {
            if (next != pre) {
                count += dfs(path, next, index, k - 1);
            }
        }
        return count;
    }
}

复杂度分析

• 时间复杂度：

  • 构建邻接表 O(n + m)

  • DFS 计算每个节点的目标节点数，第一棵树 O(nk)，第二棵树 O(mk)

  • 总体为 O((n + m) * k)

• 空间复杂度：

  • 邻接表存储 O(n + m)

  • 递归栈深度最多为 k，空间为 O(k)

总结

这道题让我意识到有时候“模拟操作”不是必须的，只要抓住最大值的来源是什么，我们就能从最优结果出发反推逻辑，像这题中只需要在第二棵树找一个最优连接点就能解决问题，极大降低了复杂度。同时，DFS 的写法也让我再次注意到在无向图中避免重复遍历的重要性，记录“上一个节点”是很有效的小技巧 ✌️

希望这篇分享能为你带来启发！如果你有任何问题或建议，欢迎在评论区留言，与我共同交流探讨。