Problem statement

https://leetcode.com/problems/distribute-coins-in-binary-tree/

Solution

This problem happens to be more complicated than I expected. Let dfs(node) will return 2 values:

  1. Current balance for given node, that is how many nodes we still need or exceed in subtree.
  2. Number of moves we need to make to make it balanced. Notice that we assume that when we do these moves, we assume that we have enough nodes in parent of node, which will be true if we use postorder traversal (no strict proof).

How we can recalculate these numbers for node if we know them for the left and for the right children?

  1. First of all ballance can be calculated as sum of balances of left and right nodes + value in node minus one.
  2. More interesting about number of moves. If we assume that temporarily we can have negative number of coint, then we just make bal number of moves to distribute coins to the left and to the right and then make balanced left and right subtrees. Why it will always work is not completely clear for me.

Complexity

It is O(n) for time and O(h) for space.

Code

class Solution:
    def distributeCoins(self, root):
        def dfs(node):
            if not node: return 0, 0
            (lbal, lcnt), (rbal, rcnt) = dfs(node.left), dfs(node.right)
            bal = node.val + lbal + rbal - 1
            return bal, lcnt + rcnt + abs(bal)
        
        return dfs(root)[1]