Distribute Coins in Binary Tree LeetCode solution in Python

1
# Definition for a binary tree node.
2
# class TreeNode:
3
#     def __init__(self, val=0, left=None, right=None):
4
#         self.val = val
5
#         self.left = left
6
#         self.right = right
7
class Solution:
8
    def distributeCoins(self, root: Optional[TreeNode]) -> int:
9

10
        def dfs(root):
11
            if not root:
12
                return 0, 0
13

14
            net_left, left_walk = dfs(root.left)
15
            net_right, right_walk = dfs(root.right)
16

17
            # if any node has extra or deficiency in both cases there has to be a walk of abs(extra) or abs(deficiency)
18

19
            return net_left + net_right + (root.val - 1), left_walk + right_walk + abs(
20
                net_left
21
            ) + abs(net_right)
22

23
        return dfs(root)[1]