Leetcode 979. Allocate coins in a binary tree

Problem: Given the root node root of a binary tree with N nodes, each node in the tree corresponds to node.val coins, and there are N coins in total.

In a move, we can select two adjacent nodes, and then move a coin from one node to another. (Movement can be from parent node to child node, or from child node to parent node.).

Returns the number of moves required to make only one coin on each node.

Source: LeetCode (LeetCode)
Link: https://leetcode-cn.com/problems/distribute-coins-in-binary-tree
Copyright is owned by LeetCode . For commercial reprints, please contact the official authorization. For non-commercial reprints, please indicate the source.

Code

class Solution {
    
    
     int ans=0;
    public int distributeCoins(TreeNode root) {
    
    
        
        dfs(root);
        return ans;
    }

    public int dfs(TreeNode node) {
    
    
        if (node == null) return 0;
        int L = dfs(node.left);
        int R = dfs(node.right);
        ans += Math.abs(L) + Math.abs(R);
        return node.val + L + R - 1;
    }
}

dfs calculates the load of the current node. When a node has a value of 0 and is a leaf node, its load is -1, when its value is n, its load is n-1; the load of non-leaf nodes is the left The load of the child and the right child is added to its own load and subtracted by 1, because a coin it has does not need to move. Then we define a member variable ans to record the total load of all nodes. The load of the root node is the number of times the coin needs to move.