leetcode —— 979. 在二叉樹中分配硬幣

給定一個有 N 個結點的二叉樹的根結點 root，樹中的每個結點上都對應有 node.val 枚硬幣，並且總共有 N 枚硬幣。

在一次移動中，我們可以選擇兩個相鄰的結點，然後將一枚硬幣從其中一個結點移動到另一個結點。(移動可以是從父結點到子結點，或者從子結點移動到父結點。)。

返回使每個結點上只有一枚硬幣所需的移動次數。

示例 1：

輸入：[3,0,0]
輸出：2
解釋：從樹的根結點開始，我們將一枚硬幣移到它的左子結點上，一枚硬幣移到它的右子結點上。

來源：力扣（LeetCode）
鏈接：https://leetcode-cn.com/problems/distribute-coins-in-binary-tree
著作權歸領釦網絡所有。商業轉載請聯繫官方授權，非商業轉載請註明出處。
————————————————
解題思路：使用深度優先遍歷，創建一個函數funs(root)，該函數的作用是計算當前結點root的金幣量（這時候當前節點的子節點的金幣數都爲1）。

當知道當前節點的左子節點的金幣數量和當前節點的右子節點的金幣數量，就可以知道當左子節點和右子節點的金幣數爲1時當前節點的金幣個數和需要移動的次數。

用變量ans保存整體需要移動的金幣次數。

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None

class Solution:
    def __init__(self):
        self.ans = 0
    def distributeCoins(self, root: TreeNode) -> int:
        def funs(root):
            if not root:
                return 0
            elif root.right and root.left:  # 如果存在左子樹和右子樹
                numl = funs(root.left) # 計算左子節點的金幣數
                numr = funs(root.right) # 計算右子節點的禁金幣數
                self.ans += abs(numl-1) + abs(numr-1) # 左子節點和右子節點的金幣數爲1需要移動的次數
                return root.val + numl + numr - 2 # 更新當前節點的金幣數
            elif root.left: # 如果右子樹爲空
                num = funs(root.left) # 計算左子節點的金幣數
                self.ans += abs(num-1)  # 左子節點金幣數爲1需要移動的次數
                return root.val + num - 1 # 更新當前節點的金幣數
            elif root.right: # 如果左子樹爲空
                num = funs(root.right)  # 計算右子節點的金幣數
                self.ans += abs(num-1)  # 右子節點金幣數爲1需要移動的次數
                return root.val + num - 1  # 更新當前節點的金幣數
            else:
                return root.val  # 如果沒有左右子樹，返回當前結點的金幣數
        funs(root) 
        return self.ans

該算法會遍歷二叉樹的每一個節點，因此時間複雜度爲O(n)，需要使用堆棧，因此空間複雜度爲O(H)，H爲二叉樹的高度。