Distribute Coins in Binary Tree

Medium · Graphs · Explanation and code example

Iterative DFS

How it works

This approach simulates recursive DFS using an explicit stack. We use a visit set to distinguish between the first visit (when we push children) and the second visit (when we process the node after its children are done). On the second visit, we accumulate coins from children, compute the extra coins, and add the absolute value to our result.

Steps

  1. Initialize a stack with the root and an empty visit set.
  2. While the stack is not empty:
    • Pop a node. If not visited, push it back, mark as visited, then push its children.
    • If already visited, add the children's values to the current node, subtract 1, and add abs(node.val) to the result.
  3. Return the total move count.

Complexity

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

Things to watch for

Forgetting to Take Absolute Value

The extra coins at a node can be positive (surplus) or negative (deficit). Both represent moves across the edge to the parent. Forgetting to take the absolute value counts deficits as negative moves, producing wrong results.

Counting Moves at the Wrong Time

In the DFS approach, moves should be counted after processing both children. Counting moves before the recursive calls or only for one child leads to incorrect totals.

Returning Wrong Value from DFS

The DFS function must return the net extra coins (surplus or deficit) that flow to the parent, not the move count. Returning the wrong value breaks the recursive accumulation.

Python code example

from typing import Optional

# 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

C++ code example

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode() : val(0), left(nullptr), right(nullptr) {}
 *     TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
 *     TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
 * };
 */
class Solution {
public:
    int distributeCoins(TreeNode* root) {
        stack<TreeNode*> stack;
        unordered_set<TreeNode*> visit;
        stack.push(root);
        int res = 0;

        while (!stack.empty()) {
            TreeNode* node = stack.top();
            stack.pop();

            if (visit.find(node) == visit.end()) {
                stack.push(node);
                visit.insert(node);

                if (node->right) {
                    stack.push(node->right);
                }
                if (node->left) {
                    stack.push(node->left);
                }
            } else {
                if (node->left) {
                    node->val += node->left->val;
                }
                if (node->right) {
                    node->val += node->right->val;
                }

                visit.erase(node);
                node->val -= 1;
                res += abs(node->val);
            }
        }

        return res;
    }
};
Source and license

Adapted from NeetCode’s source · MIT license. Node types are supplied by LeetCode where needed.

Check the explanation on LeetCode

Your recall check is in progress. Write your attempt before revealing the explanation.

Resume recall check

Browse all problem explanations