LeetCode 2641. Cousins in Binary Tree II

https://leetcode.com/problems/cousins-in-binary-tree-ii/description/

題目大意

把二元樹的每個節點值換成其 cousins 的總和
對於 node_i ， cousins 即同個 level 但與 node_i 之 parent 不同的節點們

思考

我們的想法其實很簡單，所求不就相當於把某節點的值換成 levelSum 再減掉自己跟自己的兄弟嗎？

我這邊採用 dfs 的方式解題
第一次 dfs 先求出各個 level 的 levelSum，第二次 dfs 我們再執行替換節點值成所求

class Solution
{
public:
    TreeNode *replaceValueInTree(TreeNode *root)
    {
        memset(levelSums, 0, sizeof(levelSums));
        getLevelSum(root, 0);
        root->val = 0;
        modify(root, 0);
        return root;
    }

private:
    int levelSums[100001];

    void getLevelSum(TreeNode *root, int level)
    {
        if (!root)
            return;

        levelSums[level] += root->val;
        getLevelSum(root->left, level + 1);
        getLevelSum(root->right, level + 1);
    }

    void modify(TreeNode *root, int level)
    {
        const int lVal = root->left ? root->left->val : 0;
        const int rVal = root->right ? root->right->val : 0;
        const int sibSum = lVal + rVal;
        ++level;

        if (root->left)
        {
            root->left->val = levelSums[level] - sibSum;
            modify(root->left, level);
        }
        if (root->right)
        {
            root->right->val = levelSums[level] - sibSum;
            modify(root->right, level);
        }
    }
};

Go:

func replaceValueInTree(root *TreeNode) *TreeNode {
	levelSums := make(map[int]int)
	getLevelSum(root, 0, levelSums)
	root.Val = 0
	modify(root, 0, levelSums)
	return root
}

func getLevelSum(root *TreeNode, level int, levelSums map[int]int) {
	if root == nil {
		return
	}

	levelSums[level] += root.Val
	getLevelSum(root.Left, level+1, levelSums)
	getLevelSum(root.Right, level+1, levelSums)
}

func modify(root *TreeNode, level int, levelSums map[int]int) {
	if root == nil {
		return
	}

	var lVal, rVal int
	if root.Left != nil {
		lVal = root.Left.Val
	}
	if root.Right != nil {
		rVal = root.Right.Val
	}
	siblingSum := lVal + rVal
	level++

	if root.Left != nil {
		root.Left.Val = levelSums[level] - siblingSum
		modify(root.Left, level, levelSums)
	}
	if root.Right != nil {
		root.Right.Val = levelSums[level] - siblingSum
		modify(root.Right, level, levelSums)
	}
}