🔗 🟡 979. Distribute Coins in Binary Tree 1709

tags: `Weekly Contest 120` `DFS` `貢獻法` `Binary Tree`

題意

給定一棵二元樹的根節點 $root$ ，樹上總共有 $n$ 個節點，每個節點 $node$ 上都有 $node.val$ 枚硬幣，且整棵樹上總共有 $n$ 枚硬幣。

在每次移動中，我們可以選擇相鄰的兩個節點，將一枚硬幣從其中一個節點移動到另一個節點。移動可以是從父節點到子節點，也可以是從子節點移動到父節點。

返回使每個節點上都恰好有一枚硬幣所需的最少移動次數。

思路：深度優先搜索(DFS) + 貢獻法

方法一：統計每個子樹需要和擁有的硬幣數量

對於任何一個節點，其需要的硬幣數量 $need$ 為該子樹上的節點數量，而其擁有的硬幣數量 $have$ 為其節點上的硬幣數量之和。而若 $need > have$ ，則需要從其父節點將缺少的硬幣移動到該節點上，反之則需要將多餘的硬幣移動到其父節點，故該節點對答案的貢獻為 $|need - have|$ 。

而一個樹所需要的硬幣數量 $need$ 和 $have$ 可以透過遞迴的方式計算，其中 $need$ 為左子樹和右子樹的 $need$ 之和再加上 $1$ ， $have$ 為左子樹和右子樹的 $have$ 之和再加上該節點的硬幣數量。

故可以透過 深度優先搜索(DFS) 的方式計算每個節點的 $need$ 和 $have$ ，並在DFS的過程中計算答案。

複雜度分析

時間複雜度： $O(n)$ ，其中 $n$ 為樹上的節點數量。
空間複雜度： $O(n)$ ，遞迴需要的堆疊空間。

class Solution:
    def distributeCoins(self, root: Optional[TreeNode]) -> int:
        ans = 0
        def dfs(root: TreeNode) -> Tuple[int, int]: # (need, have)
            if not root:
                return 0, 0
            nonlocal ans
            p, q = 1, root.val # need, have
            p1, q1 = dfs(root.left)
            p2, q2 = dfs(root.right)
            p += p1 + p2
            q += q1 + q2
            ans += abs(p - q)
            return p, q
        dfs(root)
        return ans

class Solution {
public:
    int ans;
    pair<int, int> dfs(TreeNode* root) {
        if (root == nullptr) return {0, 0};
        pair<int, int> res = {1, root->val}; // (need, have)
        pair<int, int> left = dfs(root->left), right = dfs(root->right);
        res.first += left.first + right.first;
        res.second += left.second + right.second;
        ans += abs(res.first - res.second); 
        return res;
    }
    int distributeCoins(TreeNode* root) {
        ans = 0;
        dfs(root);
        return ans;
    }
};

方法二：方法一優化

在方法一中， $dfs$ 函數透過返回一個 tuple 或 pair $(need, have)$ 來計算每個節點的 $need$ 和 $have$ ，但我們其實只關注兩者的差值，故可以進一步優化。

將 $dfs$ 函數改為只返回一個整數，若為正數則表示多餘的硬幣數量，負數表示缺少的硬幣數量。若只考慮該節點，不考慮左右子樹，返回值為 $root.val- 1$ ，考慮左右子樹後，返回值為 $root.val + left + right - 1$ ，且對答案的貢獻為 $|left| + |right|$ 。

複雜度分析

時間複雜度： $O(n)$ ，其中 $n$ 為樹上的節點數量。
空間複雜度： $O(n)$ ，遞迴需要的堆疊空間。

class Solution:
    def distributeCoins(self, root: Optional[TreeNode]) -> int:
        ans = 0
        def dfs(root: TreeNode) -> int: # 正數表示多餘的硬幣，負數表示缺少的硬幣
            nonlocal ans
            if not root:
                return 0
            left = dfs(root.left)
            right = dfs(root.right)
            ans += abs(left) + abs(right)
            return root.val + left + right - 1
        dfs(root)
        return ans

class Solution {
public:
    int ans;
    int dfs(TreeNode* root) { // 正數表示多餘的硬幣，負數表示缺少的硬幣
        if (root == nullptr) return 0;
        int left = dfs(root->left), right = dfs(root->right);
        ans += abs(left) + abs(right); // 要往上拿回/往下給多少硬幣
        return root->val + left + right - 1;
    }
    int distributeCoins(TreeNode* root) {
        ans = 0;
        dfs(root);
        return ans;
    }
};