2023.7.14 在二叉树中分配硬币

发布时间 2023-07-14 21:15:32作者: 烤肉kr

image

借用灵神的图:
image

所以一个直观的想法就是,计算这个子树的硬币个数和节点个数的差的绝对值。这个就是需要传递给父结点的硬币数量。
如果这个子树中的子结点也全部执行了这个操作,那么就会把硬币全部集中到当前结点,因此不用考虑子结点中的移动。所以这是个递归算法。(因为要先完成子结点的)

class Solution 
{
public:
    using PII = pair<int, int>;

    int res = 0;
    PII dfs(TreeNode *u)
    {
        if (u) 
        {
            auto l = dfs(u->left), r = dfs(u->right);
            PII data = {l.first + r.first, l.second + r.second};
            int coins = data.first + u->val;
            int nodes = data.second + 1;

            res += abs(coins - nodes);
            return {coins, nodes};
        }
        return {0, 0};
    }

    int distributeCoins(TreeNode* root) 
    {
        dfs(root);
        return res;
    }
};