LC979. Distribute Coins in Binary Tree

Approach 1: Multiple return values

# 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:
        self.res = 0
 
        def dfs(node):
            if not node:
                return (0, 0) # (size of subtree, total coins)
            
            lsize, lcoins = dfs(node.left)
            rsize, rcoins = dfs(node.right)
 
            size = 1 + lsize + rsize
            coins = node.val + lcoins + rcoins
 
            # At each point, either we have num coins = size of subtree (incl. self)
            # or we might have surplus or we might have defecit. We have to move the
            # surplus/deficit up or down the tree at each point (since only one coin
            # can be moved at a time b/w adjacent nodes).
            #
            # abs(size - coins) can give us that surplus/deficit that we need to
            # account for.
            self.res += abs(size - coins)
 
            return [size, coins]
 
        dfs(root)
 
        return self.res

Complexity

Time: $O(n)$
Space: $O(n)$

Notes

Approach 2: Single return value

# 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:
        self.res = 0
 
        def dfs(node):
            if not node:
                return 0
            
            limbalance = dfs(node.left)
            rimbalance = dfs(node.right)
 
            # One coin should be counted for current node and remaining
            # (if any) are imbalance. This could be -ve indicating we are
            # in deficit for current node and probably subtree.
            imbalance = node.val - 1 + limbalance + rimbalance
            
            # Either we need to move `imbalance` coins down to our subtree
            # if we are in deficit or up the parent if we have surplus
            self.res += abs(imbalance)
 
            return imbalance
 
        dfs(root)
        
        return self.res

Complexity

Time: $O(n)$
Space: $O(n)$

Other Languages