Hard
题目描述
给你一棵有 n 个节点的无向树,节点从 0 到 n - 1 编号,以节点 0 为根。给你一个长度为 n - 1 的二维整数数组 edges,其中 edges[i] = [ai, bi] 表示树中节点 ai 和 bi 之间有一条边。
还给你一个长度为 n 的整数数组 cost,其中 cost[i] 是分配给第 i 个节点的成本。
你需要在树的每个节点上放置一些硬币。放置在节点 i 上的硬币数量可以按以下方式计算:
- 如果节点 i 的子树大小小于 3,放置 1 个硬币。
- 否则,放置等于节点 i 子树中 3 个不同节点成本值的最大乘积的硬币数量。如果这个乘积是负数,放置 0 个硬币。
返回一个大小为 n 的数组 coin,其中 coin[i] 是放置在节点 i 上的硬币数量。
示例 1:
输入:edges = [[0,1],[0,2],[0,3],[0,4],[0,5]], cost = [1,2,3,4,5,6]
输出:[120,1,1,1,1,1]
解释:对于节点 0,放置 6 * 5 * 4 = 120 个硬币。所有其他节点都是大小为 1 的叶子节点,在每个节点上放置 1 个硬币。
示例 2:
输入:edges = [[0,1],[0,2],[1,3],[1,4],[1,5],[2,6],[2,7],[2,8]], cost = [1,4,2,3,5,7,8,-4,2]
输出:[280,140,32,1,1,1,1,1,1]
解释:放置在每个节点上的硬币数量为:
- 在节点 0 上放置 8 * 7 * 5 = 280 个硬币。
- 在节点 1 上放置 7 * 5 * 4 = 140 个硬币。
- 在节点 2 上放置 8 * 2 * 2 = 32 个硬币。
- 所有其他节点都是大小为 1 的叶子节点,在每个节点上放置 1 个硬币。
示例 3:
输入:edges = [[0,1],[0,2]], cost = [1,2,-2]
输出:[0,1,1]
解释:节点 1 和 2 是大小为 1 的叶子节点,在每个节点上放置 1 个硬币。对于节点 0,唯一可能的成本乘积是 2 * 1 * (-2) = -4。因此在节点 0 上放置 0 个硬币。
约束条件:
- 2 <= n <= 2 * 10^4
- edges.length == n - 1
- edges[i].length == 2
- 0 <= ai, bi < n
- cost.length == n
- 1 <= |cost[i]| <= 10^4
- 输入保证 edges 表示一棵有效的树。
解题思路
这道题需要使用深度优先搜索(DFS)来遍历树,并为每个节点计算其子树中3个节点的最大乘积。
核心思路:
树形DP思想:使用DFS遍历树,对于每个节点,需要收集其子树中的所有成本值,然后计算三个数的最大乘积。
关键观察:要得到三个数的最大乘积,只需要保留:
- 最大的3个正数(或0)
- 最小的3个负数(或0)
因为最大乘积只可能来自于:
- 三个最大正数的乘积
- 两个最小负数与一个最大正数的乘积
优化存储:每个子树只需要保存最多6个值(3个最大值 + 3个最小值),这样可以有效控制空间复杂度。
计算逻辑:
- 如果子树大小小于3,结果为1
- 否则对所有值排序,计算
max(最小值×次小值×最大值, 最大值×次大值×第三大值) - 如果结果为负数,返回0
算法流程:
- 构建邻接表表示树
- DFS遍历,从叶子节点开始收集成本值
- 对于每个节点,合并子节点返回的值并保留关键的6个值
- 计算当前节点的硬币数量
- 返回处理后的值供父节点使用
代码实现
class Solution {
public:
vector<long long> placedCoins(vector<vector<int>>& edges, vector<int>& cost) {
int n = cost.size();
vector<vector<int>> graph(n);
for (auto& edge : edges) {
graph[edge[0]].push_back(edge[1]);
graph[edge[1]].push_back(edge[0]);
}
vector<long long> result(n);
function<vector<int>(int, int)> dfs = [&](int node, int parent) -> vector<int> {
vector<int> values = {cost[node]};
for (int child : graph[node]) {
if (child != parent) {
vector<int> childValues = dfs(child, node);
values.insert(values.end(), childValues.begin(), childValues.end());
}
}
if (values.size() < 3) {
result[node] = 1;
return values;
}
sort(values.begin(), values.end());
int size = values.size();
long long prod1 = (long long)values[0] * values[1] * values[size - 1];
long long prod2 = (long long)values[size - 1] * values[size - 2] * values[size - 3];
long long maxProd = max(prod1, prod2);
result[node] = maxProd < 0 ? 0 : maxProd;
// Keep only the 3 smallest and 3 largest values
vector<int> keep;
for (int i = 0; i < min(3, size); i++) {
keep.push_back(values[i]);
}
for (int i = max(3, size - 3); i < size; i++) {
if (find(keep.begin(), keep.end(), values[i]) == keep.end()) {
keep.push_back(values[i]);
}
}
return keep;
};
dfs(0, -1);
return result;
}
};
class Solution:
def placedCoins(self, edges: List[List[int]], cost: List[int]) -> List[int]:
n = len(cost)
graph = [[] for _ in range(n)]
for a, b in edges:
graph[a].append(b)
graph[b].append(a)
result = [0] * n
def dfs(node, parent):
values = [cost[node]]
for child in graph[node]:
if child != parent:
child_values = dfs(child, node)
values.extend(child_values)
if len(values) < 3:
result[node] = 1
return values
values.sort()
size = len(values)
prod1 = values[0] * values[1] * values[size - 1]
prod2 = values[size - 1] * values[size - 2] * values[size - 3]
max_prod = max(prod1, prod2)
result[node] = max_prod if max_prod >= 0 else 0
# Keep only the 3 smallest and 3 largest values
keep = values[:3]
for val in values[max(3, size - 3):]:
if val not in keep:
keep.append(val)
return keep
dfs(0, -1)
return result
public class Solution {
public long[] PlacedCoins(int[][] edges, int[] cost) {
int n = cost.Length;
var graph = new List<int>[n];
for (int i = 0; i < n; i++) {
graph[i] = new List<int>();
}
foreach (var edge in edges) {
graph[edge[0]].Add(edge[1]);
graph[edge[1]].Add(edge[0]);
}
var result = new long[n];
List<int> Dfs(int node, int parent) {
var values = new List<int> { cost[node] };
foreach (int child in graph[node]) {
if (child != parent) {
var childValues = Dfs(child, node);
values.AddRange(childValues);
}
}
if (values.Count < 3) {
result[node] = 1;
return values;
}
values.Sort();
int size = values.Count;
long prod1 = (long)values[0] * values[1] * values[size - 1];
long prod2 = (long)values[size - 1] * values[size - 2] * values[size - 3];
long maxProd = Math.Max(prod1, prod2);
result[node] = maxProd < 0 ? 0 : maxProd;
// Keep only the 3 smallest and 3 largest values
var keep = new List<int>();
for (int i = 0; i < Math.Min(3, size); i++) {
keep.Add(values[i]);
}
for (int i = Math.Max(3, size - 3); i < size; i++) {
if (!keep.Contains(values[i])) {
keep.Add(values[i]);
}
}
return keep;
}
Dfs(0, -1);
return result;
}
}
var placedCoins = function(edges, cost) {
const n = cost.length;
const graph = Array(n).fill().map(() => []);
for (const [a, b] of edges) {
graph[a].push(b);
graph[b].push(a);
}
const result = new Array(n);
function dfs(node, parent) {
const values = [cost[node]];
for (const child of graph[node]) {
if (child !== parent) {
const childValues = dfs(child, node);
values.push(...childValues);
}
}
if (values.length < 3) {
result[node] = 1;
return values;
}
values.sort((a, b) => a - b);
const size = values.length;
const prod1 = values[0] * values[1] * values[size - 1];
const prod2 = values[size - 1] * values[size - 2] * values[size - 3];
const maxProd = Math.max(prod1, prod2);
result[node] = maxProd < 0 ? 0 : maxProd;
// Keep only the 3 smallest and 3 largest values
const keep = values.slice(0, 3);
for (let i = Math.max(3, size - 3); i < size; i++) {
if (!keep.includes(values[i])) {
keep.push(values[i]);
}
}
return keep;
}
dfs(0, -1);
return result;
};
复杂度分析
| 复杂度类型 | 分析 |
|---|---|
| 时间复杂度 | O(n log n),每个节点最多处理6个值,排序需要O(log 6) = O(1),总体为O(n) |
| 空间复杂度 | O(n),递归栈深度最大为树的高度,邻接表存储需要O(n)空间 |