2641 - Cousins in Binary Tree II

Given the root of a binary tree, replace the value of each node in the tree with the sum of all its cousins' values.

Two nodes of a binary tree are cousins if they have the same depth with different parents.

Return the root of the modified tree.

Note that the depth of a node is the number of edges in the path from the root node to it.

 Example 1:

Input: root = [5,4,9,1,10,null,7]
Output: [0,0,0,7,7,null,11]
Explanation: The diagram above shows the initial binary tree and the binary tree after changing the value of each node.
- Node with value 5 does not have any cousins so its sum is 0.
- Node with value 4 does not have any cousins so its sum is 0.
- Node with value 9 does not have any cousins so its sum is 0.
- Node with value 1 has a cousin with value 7 so its sum is 7.
- Node with value 10 has a cousin with value 7 so its sum is 7.
- Node with value 7 has cousins with values 1 and 10 so its sum is 11.

Example 2:

Input: root = [3,1,2]
Output: [0,0,0]
Explanation: The diagram above shows the initial binary tree and the binary tree after changing the value of each node.
- Node with value 3 does not have any cousins so its sum is 0.
- Node with value 1 does not have any cousins so its sum is 0.
- Node with value 2 does not have any cousins so its sum is 0.

 Constraints:

• The number of nodes in the tree is in the range [1, 105].
• 1 <= Node.val <= 104

C++

#include <vector>
using namespace std;

struct TreeNode {
    int val;
    TreeNode* left;
    TreeNode* right;
    TreeNode(int x) : val(x), left(NULL), right(NULL) {}
};

class Solution {
private:
    vector<int> s;

public:
    TreeNode* replaceValueInTree(TreeNode* root) {
        dfs1(root, 0);
        root->val = 0;
        dfs2(root, 1);
        return root;
    }

private:
    void dfs1(TreeNode* root, int d) {
        if (root == nullptr) {
            return;
        }
        if (s.size() <= d) {
            s.push_back(0);
        }
        s[d] += root->val;
        dfs1(root->left, d + 1);
        dfs1(root->right, d + 1);
    }

    void dfs2(TreeNode* root, int d) {
        if (root == nullptr) {
            return;
        }
        int l = (root->left == nullptr) ? 0 : root->left->val;
        int r = (root->right == nullptr) ? 0 : root->right->val;
        if (root->left != nullptr) {
            root->left->val = s[d] - l - r;
        }
        if (root->right != nullptr) {
            root->right->val = s[d] - l - r;
        }
        dfs2(root->left, d + 1);
        dfs2(root->right, d + 1);
    }
};

PYTHON

class TreeNode:
    def __init__(self, x):
        self.val = x
        self.left = None
        self.right = None

class Solution:
    def __init__(self):
        self.s = []

    def replaceValueInTree(self, root: TreeNode) -> TreeNode:
        self.dfs1(root, 0)
        root.val = 0
        self.dfs2(root, 1)
        return root

    def dfs1(self, root: TreeNode, d: int):
        if root is None:
            return
        if len(self.s) <= d:
            self.s.append(0)
        self.s[d] += root.val
        self.dfs1(root.left, d + 1)
        self.dfs1(root.right, d + 1)

    def dfs2(self, root: TreeNode, d: int):
        if root is None:
            return
        l = 0 if root.left is None else root.left.val
        r = 0 if root.right is None else root.right.val
        if root.left is not None:
            root.left.val = self.s[d] - l - r
        if root.right is not None:
            root.right.val = self.s[d] - l - r
        self.dfs2(root.left, d + 1)
        self.dfs2(root.right, d + 1)

JAVA

class Solution {
    private List<Integer> s = new ArrayList<>();

    public TreeNode replaceValueInTree(TreeNode root) {
        dfs1(root, 0);
        root.val = 0;
        dfs2(root, 1);
        return root;
    }

    private void dfs1(TreeNode root, int d) {
        if (root == null) {
            return;
        }
        if (s.size() <= d) {
            s.add(0);
        }
        s.set(d, s.get(d) + root.val);
        dfs1(root.left, d + 1);
        dfs1(root.right, d + 1);
    }

    private void dfs2(TreeNode root, int d) {
        if (root == null) {
            return;
        }
        int l = root.left == null ? 0 : root.left.val;
        int r = root.right == null ? 0 : root.right.val;
        if (root.left != null) {
            root.left.val = s.get(d) - l - r;
        }
        if (root.right != null) {
            root.right.val = s.get(d) - l - r;
        }
        dfs2(root.left, d + 1);
        dfs2(root.right, d + 1);
    }
}