树中每个节点放置的金币数目

标签: 深度优先搜索 动态规划 排序 堆(优先队列)

难度: Hard

给你一棵 n 个节点的 无向 树,节点编号为 0 到 n - 1 ,树的根节点在节点 0 处。同时给你一个长度为 n - 1 的二维整数数组 edges ,其中 edges[i] = [ai, bi] 表示树中节点 ai 和 bi 之间有一条边。

给你一个长度为 n 下标从 0 开始的整数数组 cost ,其中 cost[i] 是第 i 个节点的 开销 。

你需要在树中每个节点都放置金币,在节点 i 处的金币数目计算方法如下:

  • 如果节点 i 对应的子树中的节点数目小于 3 ,那么放 1 个金币。
  • 否则,计算节点 i 对应的子树内 3 个不同节点的开销乘积的 最大值 ,并在节点 i 处放置对应数目的金币。如果最大乘积是 负数 ,那么放置 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 * 104
  • edges.length == n - 1
  • edges[i].length == 2
  • 0 <= ai, bi < n
  • cost.length == n
  • 1 <= |cost[i]| <= 104
  • edges 一定是一棵合法的树。

Submission

运行时间: 438 ms

内存: 35.0 MB

class Solution:
    # 贪心
    # https://leetcode.cn/problems/find-number-of-coins-to-place-in-tree-nodes/solutions/2577675/tan-xin-pythonjavacgo-by-endlesscheng-j58c/
    def placedCoins(self, edges: List[List[int]], cost: List[int]) -> List[int]:
        n = len(cost)
        g = [[] for _ in range(n)]
        for x, y in edges:
            g[x].append(y)
            g[y].append(x)

        ans = [1] * n
        def dfs(x: int, fa: int) -> List[int]:
            a = [cost[x]]
            for y in g[x]:
                if y != fa:
                    a.extend(dfs(y, x))

            a.sort()
            m = len(a)
            if m >= 3:
                ans[x] = max(a[-3] * a[-2] * a[-1], a[0] * a[1] * a[-1], 0)

            if m > 5:
                a = a[:2] + a[-3:]

            return a

        dfs(0, -1)
        return ans


Explain

这个题解采用了贪心算法和深度优先搜索(DFS)的思路。首先,构建了一个邻接表表示树,然后从根节点开始对树进行深度优先搜索。对于每个节点,我们首先计算它的子树中所有节点的开销,并将它们排序。然后,如果子树中的节点数大于等于3,我们取最大的三个开销的乘积作为该节点处放置的金币数。为了优化性能,我们只保留子树中最小的两个和最大的三个开销,因为这五个值足以计算任何子树中三个节点的最大开销乘积。

时间复杂度: O(n)

空间复杂度: O(n)

class Solution:
    def placedCoins(self, edges: List[List[int]], cost: List[int]) -> List[int]:
        n = len(cost)
        g = [[] for _ in range(n)]  # 构建邻接表
        for x, y in edges:
            g[x].append(y)
            g[y].append(x)

        ans = [1] * n  # 初始化每个节点的金币数为1
        def dfs(x: int, fa: int) -> List[int]:
            a = [cost[x]]
            for y in g[x]:
                if y != fa:
                    a.extend(dfs(y, x))  # 递归计算子树的开销

            a.sort()
            m = len(a)
            if m >= 3:
                ans[x] = max(a[-3] * a[-2] * a[-1], a[0] * a[1] * a[-1], 0)  # 计算最大乘积

            if m > 5:
                a = a[:2] + a[-3:]  # 保留最小的两个和最大的三个开销

            return a

        dfs(0, -1)  # 从根节点开始深度优先搜索
        return ans

Explore

在这个问题中,我们关注的是最大的三个数的乘积。为了找到最大的三数乘积,我们需要考虑所有正数和负数的情况。正数中最大的三个乘积最大,而负数中最小的两个和最大的一个相乘可能得到一个很大的正数乘积。因此,为了确保我们能计算出所有可能的最大乘积,我们需要保留这五个数:最小的两个和最大的三个。这样做可以显著减少在递归过程中需要传递和处理的数据量,同时保证了能够计算出最大的三数乘积。

排序操作能够帮助我们快速找到最大的三个值和最小的两个值,这是计算最大乘积的关键。通过在每次递归返回时进行排序和保留关键数据,我们有效地减少了处理的数据量,并保留了计算最大乘积所必需的信息。这样的操作减少了内存的使用,并可以提高算法的执行效率,因为我们避免了在每个节点处理过多的数据。虽然排序本身是有成本的,但因为我们限制了数据的数量,所以这个成本是可控的。

如果三个数中有奇数个负数(即一个或三个),那么它们的乘积会是负数。例如,两个正数和一个负数的乘积,或者三个负数的乘积。在这种情况下,最大乘积可能为负数,表示在这些节点放置金币不是最优解,因此选择放置0个金币以避免损失。

贪心算法在这个问题中的应用主要体现在每次选择最大的三个数的乘积或者最小的两个数和最大的一个数的乘积来作为最终的金币数。这种方法假设通过局部最优选择(即在每个节点选择可能的最大乘积)可以导致全局最优解。这种策略允许算法在每一步都做出看似最优的决策,简化了问题解决的复杂性,并在很多情况下能够得到非常接近或实际的最优解。