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个节点的最大乘积。

核心思路:

  1. 树形DP思想:使用DFS遍历树,对于每个节点,需要收集其子树中的所有成本值,然后计算三个数的最大乘积。

  2. 关键观察:要得到三个数的最大乘积,只需要保留:

    • 最大的3个正数(或0)
    • 最小的3个负数(或0)

    因为最大乘积只可能来自于:

    • 三个最大正数的乘积
    • 两个最小负数与一个最大正数的乘积
  3. 优化存储:每个子树只需要保存最多6个值(3个最大值 + 3个最小值),这样可以有效控制空间复杂度。

  4. 计算逻辑

    • 如果子树大小小于3,结果为1
    • 否则对所有值排序,计算 max(最小值×次小值×最大值, 最大值×次大值×第三大值)
    • 如果结果为负数,返回0

算法流程:

  1. 构建邻接表表示树
  2. DFS遍历,从叶子节点开始收集成本值
  3. 对于每个节点,合并子节点返回的值并保留关键的6个值
  4. 计算当前节点的硬币数量
  5. 返回处理后的值供父节点使用

代码实现

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)空间

相关题目