[LeetCode] 每日一题 3373. 连接两棵树后最大目标节点数目 II（DSP + 染色）

题目描述

有两棵 无向 树，分别有 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 之间有一条边。

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

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

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

题目链接

https://leetcode.cn/problems/maximize-the-number-of-target-nodes-after-connecting-trees-ii

示例输入

示例 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]]

输出：[8,7,7,8,8]

解释：

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

示例 2

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

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

解释：

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

提示

• 2 <= n, m <= 10^5

• 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 都表示合法的树。

解题思路

这题一上来给人一种“和昨天那题很像”的感觉，只不过目标节点的定义变成了路径边数为偶数。我第一反应是直接改写昨天的思路，把判断条件换一下，谁知直接超时了。这才注意到，这题的节点数量扩大到了 1e5，暴力 DFS 已经不再适用了。

后来在讨论区找到一种特别适合“奇偶路径判断”的方法：染色法。因为树中任意两点之间的距离可以通过“节点深度差”的奇偶性来判断，我们可以给树进行二分染色（例如，深度为偶数的节点染成 0，奇数的染成 1），这样一来，路径为偶数的点对，其颜色必然相同。这就把原来的路径问题转换成了颜色统计问题，大大简化了复杂度。

具体地，我们分别对两棵树进行 DFS 染色，记录颜色为 0 和 1 的节点数量。然后对于第一棵树中的每个节点 i，只需统计它的颜色，设为 c，结果就是：

• 第一棵树中颜色为 c 的节点数 +

• 第二棵树中颜色为 0 或 1 的最大值（因为我们可以连任意节点，最大化目标节点数就是选多的那一边）

这个思路简单高效，适合大数据量下的树结构问题。

代码实现

class Solution {
    public int[] maxTargetNodes(int[][] edges1, int[][] edges2) {
        int n = edges1.length + 1, m = edges2.length + 1;
        int[] colours1 = new int[n];
        int[] colours2 = new int[m];
        int[] count1 = build(edges1, colours1);
        int[] count2 = build(edges2, colours2);
        int[] ans = new int[n];
        for (int i = 0; i < n; i++) {
            ans[i] = count1[colours1[i]] + Math.max(count2[0], count2[1]);
        }
        return ans;
    }

    private int[] build(int[][] edges, int[] colours) {
        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]);
        }
        int count0 = dfs(paths, 0, -1, 0, colours);
        return new int[] {count0, n - count0};
    }

    private int dfs(List<List<Integer>> paths, int index, int pre, int distance, int[] colours) {
        colours[index] = distance % 2;
        int count = (distance + 1) % 2;
        for (Integer next : paths.get(index)) {
            if (next != pre) {
                count += dfs(paths, next, index, distance + 1, colours);
            }
        }
        return count;
    }
}

复杂度分析

• 时间复杂度：O(n + m) 🕒
  每棵树遍历一遍节点进行染色，整体是线性复杂度。

• 空间复杂度：O(n + m) 💾
  用于存储图结构与颜色数组。

总结

这题让我意识到在大规模数据下，思路优化才是关键。路径奇偶的问题其实本质是一个颜色分类问题。以前没太注意过这种染色技巧，现在觉得在树上做问题转换真的很有用。尤其是这种“距离奇偶性”类题目，染色法会是一个非常好用的套路，以后可以作为一个固定模板来处理此类题✌️

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