在二叉树中分配硬币

67 阅读1分钟

979. 在二叉树中分配硬币 - 力扣(LeetCode)

给你一个有 n 个结点的二叉树的根结点 root ,其中树中每个结点 node 都对应有 node.val 枚硬币。整棵树上一共有 n 枚硬币。

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

返回使每个结点上 只有 一枚硬币所需的 最少 移动次数。

示例 1:

输入: root = [3,0,0]
输出: 2
解释: 一枚硬币从根结点移动到左子结点,一枚硬币从根结点移动到右子结点。

示例 2:

输入: root = [0,3,0]
输出: 3
解释: 将两枚硬币从根结点的左子结点移动到根结点(两次移动)。然后,将一枚硬币从根结点移动到右子结点。

提示:

  • 树中节点的数目为 n
  • 1 <= n <= 100
  • 0 <= Node.val <= n
  • 所有 Node.val 的值之和是 n

思路

本题可以用深度优先遍历求解。设左子树有 l 个节点,val 和为 leftVal,右子树有 r 个节点,val 和为 rightVal,假设 leftVal > l, rightVal < r, 我们需要把左子树多出的 val 先移动到根节点,再移动到右子树, 左子树多出的 val 从左子树根节点移动到 root 需要移动 leftVal - l 次,右子树少的 val 从根节点移动到右子树需要 r - rightVal 次。我们用同样的方式计算左子树和右子树各个节点需要移动的次数,最后求各节点移动次数和。代码如下。

解题

/**
 * Definition for a binary tree node.
 * function TreeNode(val, left, right) {
 *     this.val = (val===undefined ? 0 : val)
 *     this.left = (left===undefined ? null : left)
 *     this.right = (right===undefined ? null : right)
 * }
 */
/**
 * @param {TreeNode} root
 * @return {number}
 */
var distributeCoins = function (root) {
  let res = 0;
  const dfs = (node) => {
    let count = 0;
    if (node.left) {
      let left = dfs(node.left);
      count += left;
      res += Math.abs(left);
    }
    if (node.right) {
      let right = dfs(node.right);
      count += right;
      res += Math.abs(right);
    }
    count += node.val - 1;
    return count;
  };
  dfs(root);
  return res;
};