[leetcode] 979. Distribute Coins in Binary Tree @ python

原题

https://leetcode.com/problems/distribute-coins-in-binary-tree/

解法

https://www.youtube.com/watch?v=zQqku1AXVF8
DFS. 构建self.ans, 初始化为0. 定义dfs函数, 返回当前节点的流量balace. 每次递归更新self.ans. 每次递归, 累加左节点和右节点的流量的绝对值到self.ans.

balance(node) = node.val -1 + balance(node.left) + balance(node.right)
self.ans += abs(left) + abs(right)

代码

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

class Solution(object):
    def distributeCoins(self, root):
        """
        :type root: TreeNode
        :rtype: int
        """
        self.ans = 0
        
        def dfs(root):
            # return the balance of the node
            if not root: return 0
            left = dfs(root.left)
            right = dfs(root.right)
            self.ans += abs(left) + abs(right)
            return root.val -1 + left + right
        
        dfs(root)
        return self.ans