在二叉树中分配硬币

标签: 深度优先搜索 二叉树

难度: Medium

给你一个有 n 个结点的二叉树的根结点 root ,其中树中每个结点 node 都对应有 node.val 枚硬币。整棵树上一共有 n 枚硬币。

在一次移动中,我们可以选择两个相邻的结点,然后将一枚硬币从其中一个结点移动到另一个结点。移动可以是从父结点到子结点,或者从子结点移动到父结点。

返回使每个结点上 只有 一枚硬币所需的 最少 移动次数。

示例 1:

输入:root = [3,0,0]
输出:2
解释:一枚硬币从根结点移动到左子结点,一枚硬币从根结点移动到右子结点。

示例 2:

输入:root = [0,3,0]
输出:3
解释:将两枚硬币从根结点的左子结点移动到根结点(两次移动)。然后,将一枚硬币从根结点移动到右子结点。

提示:

  • 树中节点的数目为 n
  • 1 <= n <= 100
  • 0 <= Node.val <= n
  • 所有 Node.val 的值之和是 n

Submission

运行时间: 22 ms

内存: 16.0 MB

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution:
    def distributeCoins(self, root: Optional[TreeNode]) -> int:
        res=0
        def dfs(root):
            nonlocal res
            if not root:return 0,0
            l=dfs(root.left)
            r=dfs(root.right)
            res+=abs(l[0]-l[1])+abs(r[0]-r[1])
            return l[0]+r[0]+root.val,l[1]+r[1]+1
        dfs(root)
        return res

Explain

该题解采用了深度优先搜索(DFS)来解决问题。在遍历每个节点时,计算从该节点的左子树和右子树移动硬币所需的最小移动次数,并累加到全局变量 res 中。为每个节点定义一个返回值 (total_coins, total_nodes),其中 total_coins 表示以当前节点为根的子树中所有节点的硬币总数,total_nodes 表示子树中节点的总数。对于每个节点,其需要的硬币数应为其节点数 (即 total_nodes),因此其余硬币应向上或向下移动,移动的次数即为 abs(total_coins - total_nodes)。通过递归地计算每个节点的这两个值,并更新 res,最终得到全树的最小移动次数。

时间复杂度: O(n)

空间复杂度: O(n)

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution:
    def distributeCoins(self, root: Optional[TreeNode]) -> int:
        res = 0
        def dfs(root):
            nonlocal res
            if not root: return 0, 0
            l = dfs(root.left)
            r = dfs(root.right)
            # 累加左右子树的硬币移动次数
            res += abs(l[0] - l[1]) + abs(r[0] - r[1])
            # 返回当前子树的总硬币数和总节点数
            return l[0] + r[0] + root.val, l[1] + r[1] + 1
        dfs(root)
        return res

Explore

在这个问题中,`total_coins`代表的是以当前节点为根的子树中所有节点的硬币总数。这包括当前节点及其所有子节点的硬币数。而`total_nodes`则代表的是以当前节点为根的子树中节点的总数,包括当前节点自身及所有子节点。这两个值的计算对于决定每个节点需要向上或向其他节点移动多少硬币至关重要,因为理想状态下每个节点应有一个硬币。

每个节点理想上应持有一个硬币,因此如果一个节点及其子树(统称为子树)拥有的硬币数`total_coins`与节点数`total_nodes`不相等,就会存在硬币过多或不足的情况。`abs(total_coins - total_nodes)`计算的是当前子树的硬币数与其节点总数的差的绝对值,这代表了为了使每个节点都有一个硬币,需要从这个节点移动到其他节点或从其他节点移动到这个节点的硬币数。因此,这个值直接反映了从当前节点及其子树中移入或移出的硬币的最小次数。

在递归处理每个节点时,我们计算左子树和右子树分别需要移动的硬币次数,即`abs(l[0] - l[1])`和`abs(r[0] - r[1])`。这两个值分别表示从左子树和右子树中平衡硬币到每个节点所需要的移动次数。通过在每次递归调用时累加这些值到一个全局变量`res`中,我们实际上是在计算整棵树中所有不平衡的部分所需的总移动次数。因此,这种累加过程确保了我们能够得到整棵树从初始状态到每个节点都拥有一个硬币的状态所需要的最小移动次数。

返回子树的总硬币数和总节点数是为了在递归的每一层能够有效地计算出当前节点及其子树的状态,从而决定需要如何移动硬币来达到平衡状态。通过持续返回这两个值,递归函数能够在上溯过程中逐层汇总信息,并计算出每一步所需的硬币移动次数。此外,这种设计模式简化了问题的处理,因为它将问题分解为可以通过递归方法解决的较小子问题,每个子问题只关注于计算和返回其子树的硬币和节点总数。这样的设计提高了算法的效率并降低了实现复杂性。