Distribute Coins in Binary Tree

Intuition

For each edge between a parent and a node, the numbers of coins to be moved across the edge is the absolute value of the number of nodes in the sub-tree whose root is the node and the number of coins in this sub-tree. So we could solve this problem easily by one single DFS.

Solution


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

class Solution:
    def distributeCoins(self, root: TreeNode) -> int:
        return self._helper(root)[2]

    def _helper(self, node):
        if node is None:
            return 0, 0, 0
        l = self._helper(node.left)
        r = self._helper(node.right)
        nodes = l[0] + r[0] + 1
        coins = l[1] + r[1] + node.val
        moves = l[2] + r[2] + abs(nodes - coins)
        return nodes, coins, moves