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

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

题目

题目解析及思路

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

目标节点：距离<=k

暴力枚举第二棵树中每个节点到其他节点距离d<=k-1的最大数量max2

再枚举第一棵树中每个节点求目标节点数量dfs(edges1)，二者相加为答案ans = max2 + dfs

代码

class Solution {
public:
    vector<vector<int>> build_tree(vector<vector<int>>& edges){
        vector<vector<int>> g(edges.size() + 1);
        for(auto& t:edges){
            //邻接矩阵
            int x = t[0],y = t[1];
            g[x].push_back(y);
            g[y].push_back(x);
        }
        return g;
    }

    int dfs(int x,int fa,int d,vector<vector<int>>& g,int k){
        if(d > k)
            return 0;
        int cnt = 1;
        //遍历
        for(int y:g[x])
            //不是父节点就继续dfs
            if(y != fa)
                cnt += dfs(y,x,d+1,g,k);
        return cnt;
    }

    vector<int> maxTargetNodes(vector<vector<int>>& edges1, vector<vector<int>>& edges2, int k) {
        int max2 = 0;
        if(k > 0){
            //建立邻接矩阵
            auto g = build_tree(edges2);
            for(int i=0;i<edges2.size()+1;i++)
                //求距离k-1的最大数目
                max2 = max(max2,dfs(i,-1,0,g,k-1));
        }

        auto g = build_tree(edges1);
        vector<int> ans(edges1.size() + 1);
        for(int i=0;i<ans.size();i++){
            ans[i] = dfs(i,-1,0,g,k) + max2;
        }
        return ans;
    }
};