Distribute Coins In Binary Tree | LeetCode

Problem Source

Given binary tree with N nodes, in which each node value represents the number of coins in that node. It is guaranteed that there are N coins in total. In one move, you can choose to transport one coin from parent to a children or vice versa.

Find the minimum number of moves which you can make, to make sure that every node has exactly one coin.

Example #1

 

Input:            3
                /   \
               0    0
Output: 2
Explanation: We can move one coin from root of tree to left child in one move and one coin from root of tree to right child. After this, the left and the right child have one coin each and the root node also has 1 coin (3 - 2) Therefore, it takes two moves to make sure that every node has exactly one coin, and the answer is 2.

Example #2

Input:             0
                 /   \
               3     0
Output:  3
Explanation: We need to move one coin from left child to root which will make one move. After that we will move one coin from left child to right child via root node. This will take 2 moves. Total there will be 3 moves.

Solution

At each node we will be performing two types of transactions :

1. Giving coins to left subtree or right subtree (if needed) like in example #1
2. Taking coins from left subtree or right subtree (if they have excess) like in example #2

Let’s say we take 2 coins from left subtree and give 3 coins to the right subtree, and the current nodes has 1 coins at the begining. In that case, the current node has a balance of 1 + 2 – 3 = 0. Therefore, the current node will have to borrow 0 – 1 = -1 coins from his parent to make sure that it has 1 coin. (It needs to store 1 coin within itself also).

We will simulate this process by starting from the deepest nodes and returning the balance to their parent nodes. To solve this problem then, we can use DFS/Recursion as shown below

Solution Implementation

/**
 * 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) {
        int ans = 0;
        dfs(root, ans);
        return ans;
    }
  
    /** Returns a pair of integers (no of nodes in subtree, balance ) */
    int dfs(TreeNode* treeNode, int& ans) {
      if(treeNode == NULL) {
        return 0;
      }
      
      int left = dfs(treeNode->left, ans);
      int right = dfs(treeNode->right, ans);
      
      /** Add the transactions done with left child and right child */
      ans += abs(left);
      ans += abs(right);
      
   /** Available balance after doing all the transactions */
      return treeNode->val + left + right - 1;
    }
};