LeetCode 3373.连接两棵树后最大目标节点数目 II：脑筋急转弯+广度优先搜索(黑白染色法)

【LetMeFly】3373.连接两棵树后最大目标节点数目 II：脑筋急转弯+广度优先搜索(黑白染色法)

力扣题目链接：https://leetcode.cn/problems/maximize-the-number-of-target-nodes-after-connecting-trees-ii/

有两棵 无向 树，分别有 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 的 目标节点 数目的 最大值 。

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

示例 1：

输出：[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：

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

解释：

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

提示：

• 2 <= n, m <= 105
• 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 都表示合法的树。

解题方法：黑白染色法

做这道题之前可以先看下3372.连接两棵树后最大目标节点数目 I：脑筋急转弯——深搜确定k邻近节点(清晰题解)的解题思路。

这道题同理，tree1和tree2可以分开来求。

将tree1和tree2的0节点染成黑色，广搜遍历完整棵树，黑色相邻的节点染成白色(反之亦然)。

统计每棵树中黑色节点和白色节点分别一共有多少个。

对于tree1中的某个节点，tree1中和它距离为偶数的节点就是和它同色的节点数，而tree2中和它距离为偶数的点可以恒定选择$max(black,white)$。

tree2中黑色节点多的话，就将tree1[i]连接到tree2的任意一个白色节点上，那么tree2的每个黑色节点到tree1[i]的距离都是偶数；

tree2中白色节点多的话，就将tree1[i]连接到tree2的任意一个黑色节点上，那么tree2的每个白色节点到tree1[i]的距离都是偶数；

不知道有没有注意题目中每棵树的节点个数都至少为2个，所以tree2一定既有黑色节点又有白色节点。

• 时间复杂度$O(m+n)$
• 空间复杂度$O(m+n)$

AC代码

C++

/*
 * @Author: LetMeFly
 * @Date: 2025-05-29 22:14:08
 * @LastEditors: LetMeFly.xyz
 * @LastEditTime: 2025-05-29 22:50:54
 */
class Solution {
private:
    vector<vector<int>> buildTree(vector<vector<int>>& edges) {
        vector<vector<int>> ans(edges.size() + 1);
        for (vector<int>& edge : edges) {
            ans[edge[0]].push_back(edge[1]);
            ans[edge[1]].push_back(edge[0]);
        }
        return ans;
    }

    vector<bool> tree2color(vector<vector<int>>& tree) {
        vector<bool> ans(tree.size());
        queue<array<int, 3>> q;  // [<lastNode, thisNode, color>, <...>, ...
        q.push({-1, 0, 0});
        while (q.size()) {
            auto [lastNode, thisNode, color] = q.front();
            q.pop();
            ans[thisNode] = color;
            for (int nextNode : tree[thisNode]) {
                if (nextNode != lastNode) {
                    q.push({thisNode, nextNode, color ? 0 : 1});
                }
            }
        }
        return ans;
    }

    pair<int, int> count01(vector<bool>& color) {
        pair<int, int> ans;
        for (bool b : color) {
            if (b) {
                ans.second++;
            } else {
                ans.first++;
            }
        }
        return ans;
    }
public:
    vector<int> maxTargetNodes(vector<vector<int>>& edges1, vector<vector<int>>& edges2) {
        vector<vector<int>> graph1 = move(buildTree(edges1)), graph2 = move(buildTree(edges2));
        vector<bool> color1 = tree2color(graph1), color2 = tree2color(graph2);
        auto [black1, white1] = count01(color1);
        auto [black2, white2] = count01(color2);
        int toAdd = max(black2, white2);
        black1 += toAdd, white1 += toAdd;
        vector<int> ans(color1.size());
        for (int i = 0; i < ans.size(); i++) {
            ans[i] = color1[i] ? white1 : black1;
        }
        return ans;
    }
};

Python

'''
Author: LetMeFly
Date: 2025-05-29 22:14:08
LastEditors: LetMeFly.xyz
LastEditTime: 2025-05-29 23:40:08
'''
from typing import List, Tuple

class Solution:
    def edges2tree(self, edges: List[List[int]]) -> List[List[int]]:
        ans = [[] for _ in range(len(edges) + 1)]
        for x, y in edges:
            ans[x].append(y)
            ans[y].append(x)
        return ans
    
    def tree2color(self, tree: List[List[int]]) -> List[bool]:
        ans = [False] * len(tree)
        q = [(-1, 0, True)]
        while q:
            lastNode, thisNode, color = q[-1]
            q = q[:-1]
            ans[thisNode] = color
            for nextNode in tree[thisNode]:
                if nextNode != lastNode:
                    q.append((thisNode, nextNode, False if color else True))
        return ans
    
    def count01(self, color: List[bool]) -> Tuple[int, int]:
        ans = [0, 0]
        for c in color:
            ans[c] += 1
        return (*ans, )
    
    def maxTargetNodes(self, edges1: List[List[int]], edges2: List[List[int]]) -> List[int]:
        graph1, graph2 = self.edges2tree(edges1), self.edges2tree(edges2)
        color1, color2 = self.tree2color(graph1), self.tree2color(graph2)
        toAdd = max(self.count01(color2))
        black1, white1 = self.count01(color1)
        black1, white1 = black1 + toAdd, white1 + toAdd
        return [white1 if c else black1 for c in color1]

同步发文于CSDN和我的个人博客，原创不易，转载经作者同意后请附上原文链接哦~

千篇源码题解已开源