Distribute Coins in Binary Tree

lcpy gen -n 979   # by problem number
lcpy gen -s distribute_coins_in_binary_tree   # by problem name

Input: root = [3,0,0]
Output: 2
Explanation: From the root of the tree, we move one coin to its left child, and one coin to its right child.

Input: root = [0,3,0]
Output: 3
Explanation: From the left child of the root, we move two coins to the root [taking two moves]. Then, we move one coin from the root of the tree to the right child.

from leetcode_py import TreeNode


class Solution:
    # Time: O(n)
    # Space: O(n)
    def distribute_coins(self, root: TreeNode[int] | None) -> int:
        moves = 0

        def dfs(node: TreeNode[int] | None) -> int:
            nonlocal moves
            if node is None:
                return 0
            left = dfs(node.left)
            right = dfs(node.right)
            moves += abs(left) + abs(right)
            return node.val + left + right - 1

        dfs(root)
        return moves