979. Distribute Coins in Binary Tree - Explanation

Problem Link



# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution:
    def distributeCoins(self, root: Optional[TreeNode]) -> int:
        self.res = 0

        def dfs(cur):
            if not cur:
                return [0, 0]  # [size, coins]

            l_size, l_coins = dfs(cur.left)
            r_size, r_coins = dfs(cur.right)

            size = 1 + l_size + r_size
            coins = cur.val + l_coins + r_coins
            self.res += abs(size - coins)

            return [size, coins]

        dfs(root)
        return self.res

Time & Space Complexity

  • Time complexity: O(n)O(n)
  • Space complexity: O(n)O(n) for recursion stack.

2. Depth First Search (Optimal)

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution:
    def distributeCoins(self, root: Optional[TreeNode]) -> int:
        self.res = 0

        def dfs(cur):
            if not cur:
                return 0  # extra_coins

            l_extra = dfs(cur.left)
            r_extra = dfs(cur.right)

            extra_coins = cur.val - 1 + l_extra + r_extra
            self.res += abs(extra_coins)
            return extra_coins

        dfs(root)
        return self.res

Time & Space Complexity

  • Time complexity: O(n)O(n)
  • Space complexity: O(n)O(n) fo recursion stack.

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution:
    def distributeCoins(self, root: Optional[TreeNode]) -> int:
        res = 0
        q = deque([root])
        parent_map = {}

        nodes = []
        while q:
            node = q.popleft()
            nodes.append(node)
            if node.left:
                parent_map[node.left] = node
                q.append(node.left)
            if node.right:
                parent_map[node.right] = node
                q.append(node.right)

        while nodes:
            node = nodes.pop()
            if node in parent_map:
                parent = parent_map[node]
                parent.val += node.val - 1
                res += abs(node.val - 1)

        return res

Time & Space Complexity

  • Time complexity: O(n)O(n)
  • Space complexity: O(n)O(n)

4. Iterative DFS

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution:
    def distributeCoins(self, root: Optional[TreeNode]) -> int:
        stack = [root]
        res = 0
        visit = set()

        while stack:
            node = stack.pop()

            if node not in visit:
                stack.append(node)
                visit.add(node)

                if node.right:
                    stack.append(node.right)
                if node.left:
                    stack.append(node.left)
            else:
                if node.left:
                    node.val += node.left.val
                if node.right:
                    node.val += node.right.val

                node.val -= 1
                res += abs(node.val)
                visit.remove(node)

        return res

Time & Space Complexity

  • Time complexity: O(n)O(n)
  • Space complexity: O(n)O(n)