题目:3373. 连接两棵树后最大目标节点数目 II
思路:贪心+深度优先搜索dfs,时间复杂度0(n+m)。
第二棵树:对每个节点进行分类,0或1,相邻的节点肯定不同啦,这样就可以统计出0和1 各自的节点个数。
对于第一颗树而言,也是一样处理的,细节看注释。
C++版本:
class Solution {
public:
// 构建链接表
void solve(vector<vector<int>> & g,vector<vector<int>>& edges){
for(auto e:edges){
int x=e[0],y=e[1];
g[x].push_back(y);
g[y].push_back(x);
}
}
// dfs第二棵树,统计0、1节点的个数
void dfs1(int u,int fa,int d,vector<vector<int>> & g,vector<int> &cnt1){
cnt1[d]++;
for(auto x:g[u]){
if(x==fa) continue;
dfs1(x,u,d^1,g,cnt1);
}
}
// dfs第一颗树,统计0、1节点的个数,同时用数组ans记录每个节点u所在的节点状态d(0,1)
void dfs2(int u,int fa,int d,vector<vector<int>> & g,vector<int> &cnt1,vector<int> &ans){
cnt1[d]++;
ans[u]=d;
for(auto x:g[u]){
if(x==fa) continue;
dfs2(x,u,d^1,g,cnt1,ans);
}
}
vector<int> maxTargetNodes(vector<vector<int>>& edges1, vector<vector<int>>& edges2) {
// 构建链接表
int n=edges1.size(),m=edges2.size();
vector<vector<int>> g1(n+1),g2(m+1);
solve(g1,edges1);
solve(g2,edges2);
// 统计0、1节点的个数
vector<int> cnt1(2,0);
// 先dfs第二棵树
dfs1(0,-1,0,g2,cnt1);
// 找出最大的类别用于和第一颗树的节点相连
// 注意m是大于等于2的,所以一定会有一个类别最大。不存在m=1的情况,导致出问题
int mx=max(cnt1[0],cnt1[1]);
// 答案
vector<int> ans(n+1,0);
// 初始化,用于记录第一颗树的0、1节点的个数
cnt1[0]=0,cnt1[1]=0;
// dfs第二棵树
dfs2(0,-1,0,g1,cnt1,ans);
for(int i=0;i<=n;i++){
ans[i]=cnt1[ans[i]]+mx;
}
return ans;
}
};
JAVA版本:
class Solution {
void solve(List<List<Integer>> g,int[][] edges){
for(var e:edges){
int x=e[0],y=e[1];
g.get(x).add(y);
g.get(y).add(x);
}
}
void dfs1(int u,int fa,int d,List<List<Integer>> g,int[] cnt1){
cnt1[d]++;
for(var x:g.get(u)){
if(x==fa) continue;
dfs1(x,u,d^1,g,cnt1);
}
}
void dfs2(int u,int fa,int d,List<List<Integer>> g,int[] cnt1,int[] ans){
cnt1[d]++;
ans[u]=d;
for(var x:g.get(u)){
if(x==fa) continue;
dfs2(x,u,d^1,g,cnt1,ans);
}
}
public int[] maxTargetNodes(int[][] edges1, int[][] edges2) {
int n=edges1.length,m=edges2.length;
List<List<Integer>> g1=new ArrayList<>();
List<List<Integer>> g2=new ArrayList<>();
for(int i=0;i<=n;i++){
g1.add(new ArrayList<>());
}
for(int i=0;i<=m;i++){
g2.add(new ArrayList<>());
}
solve(g1,edges1);
solve(g2,edges2);
int[] cnt1=new int[2];
dfs1(0,-1,0,g2,cnt1);
int mx=Math.max(cnt1[0],cnt1[1]);
int[] ans=new int[n+1];
cnt1[0]=0;
cnt1[1]=0;
dfs2(0,-1,0,g1,cnt1,ans);
for(int i=0;i<=n;i++){
ans[i]=cnt1[ans[i]]+mx;
}
return ans;
}
}
Go版本:
func maxTargetNodes(edges1 [][]int, edges2 [][]int) []int {
n:=len(edges1)
m:=len(edges2)
g1:=make([][]int,n+1)
g2:=make([][]int,m+1)
solve(edges1,g1)
solve(edges2,g2)
cnt:=make([]int,2)
dfs1(0,-1,0,g2,cnt)
mx:=max(cnt[0],cnt[1])
cnt[0],cnt[1]=0,0
ans:=make([]int,n+1)
dfs2(0,-1,0,g1,cnt,ans)
for i:=0;i<=n;i++ {
ans[i]=cnt[ans[i]]+mx
}
return ans
}
func solve(edges [][]int,g [][]int) {
for _,e:= range edges {
x,y:=e[0],e[1]
g[x]=append(g[x],y)
g[y]=append(g[y],x)
}
}
func dfs1(u int,fa int,d int,g [][]int,cnt []int){
cnt[d]++
for _,x:=range g[u] {
if x!=fa {
dfs1(x,u,d^1,g,cnt)
}
}
}
func dfs2(u int,fa int,d int,g [][]int,cnt []int,ans []int){
ans[u]=d
cnt[d]++
for _,x:=range g[u] {
if x!=fa {
dfs2(x,u,d^1,g,cnt,ans)
}
}
}