力扣3373.连接两棵树后最大目标节点数目II

力扣3373.连接两棵树后最大目标节点数目II

题目

题目解析及思路

题目要求对于第一棵树中每个节点，使其与第二棵树的一个节点连边，返回最多的目标节点数目

目标节点：距离为偶数

染色，从节点0开始递归求两棵树中黑白两色的节点数量

第二棵树中数量多的那个颜色为max2

第一棵树中同色 + max2为答案

代码

class Solution {
public:
    vector<int> maxTargetNodes(vector<vector<int>>& edges1, vector<vector<int>>& edges2) {
        // 计数辅助函数，统计图中每个节点的颜色（0 或 1）
        auto count = [](vector<vector<int>>& edges) {
            vector<vector<int>> g(edges.size() + 1);  // 创建邻接表
            // 将边插入到图的邻接表中
            for (auto& e : edges) {
                int x = e[0], y = e[1];
                g[x].push_back(y);  // 在 x 节点和 y 节点之间添加边
                g[y].push_back(x);  // 在 y 节点和 x 节点之间添加边
            }

            array<int, 2> cnt{};  // 用来统计颜色为 0 和 1 的节点数量
            // 深度优先搜索（DFS）函数，用来遍历图并计算节点的颜色
            auto dfs = [&](auto&& dfs, int x, int fa, int d) -> void {
                cnt[d]++;  // 记录当前节点颜色的数量
                for (int y : g[x]) {
                    if (y != fa) {  // 防止回到父节点
                        dfs(dfs, y, x, d ^ 1);  // 切换颜色，继续递归
                    }
                }
            };
            dfs(dfs, 0, -1, 0);  // 从节点 0 开始深度优先搜索，初始颜色为 0
            return make_pair(g, cnt);  // 返回图的邻接表和颜色计数
        };

        // 计算 edges2 图的节点颜色数量
        auto [_, cnt2] = count(edges2);
        int max2 = max(cnt2[0], cnt2[1]);  // 找到最大颜色数量

        // 计算 edges1 图的节点颜色数量
        auto [g, cnt1] = count(edges1);
        vector<int> ans(g.size(), max2);  // 初始化答案数组，大小为图的节点数，初始值为 max2

        // 深度优先搜索，计算每个节点的最大目标节点数量
        auto dfs = [&](auto&& dfs, int x, int fa, int d) -> void {
            ans[x] += cnt1[d];  // 根据颜色 d 更新当前节点的目标节点数量
            for (int y : g[x]) {
                if (y != fa) {  // 防止回到父节点
                    dfs(dfs, y, x, d ^ 1);  // 切换颜色，继续递归
                }
            }
        };
        dfs(dfs, 0, -1, 0);  // 从节点 0 开始深度优先搜索，初始颜色为 0

        return ans;  // 返回每个节点的最大目标节点数量
    }
};