千问 LeetCode 2973. 树中每个节点放置的金币数目 Go实现

以下是 LeetCode 2973. 树中每个放置的金币数目 的 Go 语言实现：

package main

import (
"sort"
)

func placedCoins(edges [][]int, cost []int) []int64 {
n := len(cost)
// 1. 建图 (邻接表)
g := make([][]int, n)
for _, edge := range edges {
u, v := edge[0], edge[1]
g[u] = append(g[u], v)
g[v] = append(g[v], u)
}

ans := make([]int64, n)
// 默认每个节点初始金币数为 1
for i := range ans {
ans[i] = 1
}

// 2. 深度优先搜索
var dfs func(node, fa int) []int
dfs = func(node, fa int) []int {
// 收集当前子树的节点索引
subtree := []int{node}

for _, neighbor := range g[node] {
if neighbor == fa {
continue // 无向树深搜防环，跳过父节点
}
sub := dfs(neighbor, node)
if len(sub) < 5 {
subtree = append(subtree, sub...)
} else {
// 只保留2个最小的和3个最大的
subtree = append(subtree, sub[0], sub[1], sub[len(sub)-3], sub[len(sub)-2], sub[len(sub)-1])
}
}

// 对当前子树的候选节点按 cost 排序
sort.Slice(subtree, func(i, j int) bool {
return cost[subtree[i]] < cost[subtree[j]]
})

// 如果节点数大于5，只保留极值
if len(subtree) > 5 {
subtree = []int{
subtree[0], subtree[1],
subtree[len(subtree)-3], subtree[len(subtree)-2], subtree[len(subtree)-1],
}
}

// 计算当前节点应放置的金币数
if len(subtree) >= 3 {
sz := len(subtree)
// 组合1: 两个最小的负数 * 最大的正数
prod1 := int64(cost[subtree[0]]) * int64(cost[subtree[1]]) * int64(cost[subtree[sz-1]])
// 组合2: 三个最大的正数
prod2 := int64(cost[subtree[sz-3]]) * int64(cost[subtree[sz-2]]) * int64(cost[subtree[sz-1]])

maxProd := prod1
if prod2 > maxProd {
maxProd = prod2
}

if maxProd > 0 {
ans[node] = maxProd
} else {
ans[node] = 0
}
}

return subtree
}

dfs(0, -1)
return ans
}

核心逻辑与 Go 语言特性说明：
1. 闭包与状态传递：在 Go 中，我们使用匿名函数 dfs 作为闭包，这样可以直接访问外部的 cost 数组和 ans 结果数组，避免了在函数参数中反复传递这些大数组，代码更加简洁。
2. 数学规律：要在一个包含正负数的数组中选出3个数使乘积最大，只有两种可能：
- 三个最大的正数
- 一个最大的正数和两个最小的负数
3. 剪枝优化：基于上述规律，我们在深度优先搜索（DFS）合并子树时，不需要保留所有的节点。每个子树最多只需要保留 2个最小的 和 3个最大的 节点（共5个）即可。这大幅降低了存储和排序的复杂度，使得时间复杂度维持在 O(N log N) 级别。
4. 防环处理：由于是无向树，DFS 时需要传入父节点 fa，遇到父节点直接跳过，防止走回头路形成死循环。
5. 结果计算：当子树节点数不足3个时，金币数保持初始值1；当大于等于3个时，取上述两种组合乘积的最大值，若为负数则放置0个金币。注意 Go 中 int 的乘法可能会溢出，计算乘积时使用了 int64 进行强转。