Leetcode 2973. Find Number of Coins to Place in Tree Nodes

• Leetcode 2973. Find Number of Coins to Place in Tree Nodes
  • 1. 解题思路
  • 2. 代码实现

• 题目链接：2973. Find Number of Coins to Place in Tree Nodes

1. 解题思路

这道题思路上其实挺简单的，就是一个遍历的思路，找到每一个点对应的子树当中所有的节点，然后按照条件进行赋值即可。

不过，直接地实现会导致超时问题的问题，因此我们对此需要做一下剪枝，具体来说的话，由于我们要求取3个元素的最大乘积，因此考虑到正负性，选择上必然只有两种情况：

1. 最大的三个元素
2. 最大的一个元素与最小的两个元素

因此，我们事实上不需要保留全部的元素，只需要排序之后对每一个子树保留至多5个元素即可，从而大幅简化我们的存储还有排序复杂度。

2. 代码实现

给出python代码实现如下：

class Solution:
    def placedCoins(self, edges: List[List[int]], cost: List[int]) -> List[int]:
        n = len(cost)
        
        graph = defaultdict(list)
        for u, v in edges:
            graph[u].append(v)
            graph[v].append(u)
        
        tree = {}
        def dfs(root, parent):
            nonlocal tree
            
            subtree = [root]
            for node in graph[root]:
                if node == parent:
                    continue
                sub = dfs(node, root)
                if len(sub) < 5:
                    subtree.extend(sub)
                else:
                    subtree.extend(sub[:2] + sub[-3:])
            subtree = sorted(subtree, key=lambda x: cost[x])
            tree[root] = subtree
            return subtree
        
        dfs(0, -1)
        ans = [1 for _ in range(n)]
        for i in range(n):
            subtree = tree[i]
            if len(subtree) < 3:
                continue
            ans[i] = max(
                0, 
                cost[subtree[0]] * cost[subtree[1]] * cost[subtree[-1]], 
                cost[subtree[-1]] * cost[subtree[-2]] * cost[subtree[-3]]
            )
        return ans

提交代码评测得到：耗时1851ms，占用内存38.5MB。