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.
# 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