以下是 LeetCode 3373. 连接两棵树后最大目标节点数目 II 的 Go 实现:

解题思路

本题核心思路是 黑白染色法(二分图染色):

1. 关键观察:树中任意两节点间路径边数为偶数,当且仅当它们的深度奇偶性相同。我们可以将树进行二分染色(相邻节点颜色不同),则同色节点之间的距离一定是偶数。

2. 问题转化:对于第一棵树的节点 `i`,它的目标节点数 = 第一棵树中与它同色的节点数 + 第二棵树中能通过连接边变为距离偶数的最大节点数。

3. 第二棵树的贡献:由于我们可以自由选择连接第二棵树的任意节点,若节点 `i` 颜色为 `c`,我们只需连接第二棵树中颜色为 `c^1` 的节点,这样第二棵树中所有颜色为 `c` 的节点到 `i` 的距离都是偶数。因此第二棵树的最大贡献为 `max(cnt2[0], cnt2[1])`。

4. 第一棵树的贡献:`cnt1[c1[i]]` 即第一棵树中与节点 `i` 同色的节点数。

Go 代码

```go
func maxTargetNodes(edges1 [][]int, edges2 [][]int) []int {
    g1 := build(edges1)
    g2 := build(edges2)
    n, m := len(g1), len(g2)
    c1 := make([]int, n)
    c2 := make([]int, m)
    cnt1 := make([]int, 2)
    cnt2 := make([]int, 2)

    // 先对第二棵树DFS,统计两种颜色的节点数
    dfs(g2, 0, -1, c2, 0, cnt2)
    // 再对第一棵树DFS,统计两种颜色的节点数,同时记录每个节点的颜色
    dfs(g1, 0, -1, c1, 0, cnt1)

    // 第二棵树能贡献的最大目标节点数
    t := max(cnt2[0], cnt2[1])
    ans := make([]int, n)
    for i := 0; i < n; i++ {
        // 第一棵树中节点i的目标节点数 = 第二棵树最大贡献 + 第一棵树中与i同色的节点数
        ans[i] = t + cnt1[c1[i]]
    }
    return ans
}

// 根据边列表构建邻接表
func build(edges [][]int) [][]int {
    n := len(edges) + 1
    g := make([][]int, n)
    for _, e := range edges {
        a, b := e[0], e[1]
        g[a] = append(g[a], b)
        g[b] = append(g[b], a)
    }
    return g
}

// DFS遍历树,进行黑白染色
// a: 当前节点, fa: 父节点, c: 颜色数组, d: 当前颜色(0或1), cnt: 颜色计数
func dfs(g [][]int, a, fa int, c []int, d int, cnt []int) {
    c[a] = d
    cnt[d]++
    for _, b := range g[a] {
        if b != fa {
            // 子节点颜色与父节点相反 (d ^ 1)
            dfs(g, b, a, c, d^1, cnt)
        }
    }
}
```

复杂度分析

- 时间复杂度:O(n + m),每棵树只需一次 DFS 遍历
- 空间复杂度:O(n + m),邻接表和颜色数组

下载文件:[leetcode_3373.go](sandbox:///mnt/agents/output/leetcode_3373.gk

 

更多推荐