979. 在二叉树中分配硬币

59 阅读1分钟

题目:
给定一个有 N 个结点的二叉树的根结点 root,树中的每个结点上都对应有 node.val 枚硬币,并且总共有 N 枚硬币。

在一次移动中,我们可以选择两个相邻的结点,然后将一枚硬币从其中一个结点移动到另一个结点。(移动可以是从父结点到子结点,或者从子结点移动到父结点。)。

返回使每个结点上只有一枚硬币所需的移动次数。
算法:
方法一:DFS
把移入移出都当作移出

/**
 * Definition for a binary tree node.
 * type TreeNode struct {
 *     Val int
 *     Left *TreeNode
 *     Right *TreeNode
 * }
 */
func distributeCoins(root *TreeNode) int {
    ans := 0
    var dfs func(node *TreeNode) int
    dfs = func(node *TreeNode) int {
        if node == nil {
            return 0
        }
        l := dfs(node.Left)
        ans = ans + abs(l)
        r := dfs(node.Right)
        ans = ans +   abs(r)
        // 左节点从root移出+右节点从root移出+当前节点数-当前节点保留一个
        return l + r + node.Val - 1
    }
    // if dfs(root) != 0 {
    //     return - 1
    // }
    dfs(root)
    return ans
}

func abs(a int) int {
    if a < 0 {
        return -a
    }
    return a
}