Skip to content

2641. Cousins in Binary Tree II 👍

  • Time: $O(n)$
  • Space: $O(n)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
class Solution {
 public:
  TreeNode* replaceValueInTree(TreeNode* root) {
    vector<int> levelSums;
    dfs(root, 0, levelSums);
    return replace(root, 0, new TreeNode(0), levelSums);
  }

 private:
  void dfs(TreeNode* root, int level, vector<int>& levelSums) {
    if (root == nullptr)
      return;
    if (levelSums.size() == level)
      levelSums.push_back(0);
    levelSums[level] += root->val;
    dfs(root->left, level + 1, levelSums);
    dfs(root->right, level + 1, levelSums);
  }

  TreeNode* replace(TreeNode* root, int level, TreeNode* curr,
                    const vector<int>& levelSums) {
    const int nextLevel = level + 1;
    const int nextLevelCousinsSum =
        nextLevel >= levelSums.size()
            ? 0
            : levelSums[nextLevel] -
                  (root->left == nullptr ? 0 : root->left->val) -
                  (root->right == nullptr ? 0 : root->right->val);
    if (root->left != nullptr) {
      curr->left = new TreeNode(nextLevelCousinsSum);
      replace(root->left, level + 1, curr->left, levelSums);
    }
    if (root->right != nullptr) {
      curr->right = new TreeNode(nextLevelCousinsSum);
      replace(root->right, level + 1, curr->right, levelSums);
    }
    return curr;
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
class Solution {
  public TreeNode replaceValueInTree(TreeNode root) {
    List<Integer> levelSums = new ArrayList<>();
    dfs(root, 0, levelSums);
    return replace(root, 0, new TreeNode(0), levelSums);
  }

  private void dfs(TreeNode root, int level, List<Integer> levelSums) {
    if (root == null)
      return;
    if (levelSums.size() == level)
      levelSums.add(0);
    levelSums.set(level, levelSums.get(level) + root.val);
    dfs(root.left, level + 1, levelSums);
    dfs(root.right, level + 1, levelSums);
  }

  private TreeNode replace(TreeNode root, int level, TreeNode curr, List<Integer> levelSums) {
    final int nextLevel = level + 1;
    final int nextLevelCousinsSum = nextLevel >= levelSums.size()
                                        ? 0
                                        : levelSums.get(nextLevel) -
                                              (root.left == null ? 0 : root.left.val) -
                                              (root.right == null ? 0 : root.right.val);
    if (root.left != null) {
      curr.left = new TreeNode(nextLevelCousinsSum);
      replace(root.left, level + 1, curr.left, levelSums);
    }
    if (root.right != null) {
      curr.right = new TreeNode(nextLevelCousinsSum);
      replace(root.right, level + 1, curr.right, levelSums);
    }
    return curr;
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
class Solution:
  def replaceValueInTree(self, root: Optional[TreeNode]) -> Optional[TreeNode]:
    levelSums = []

    def dfs(root: Optional[TreeNode], level: int) -> None:
      if not root:
        return
      if len(levelSums) == level:
        levelSums.append(0)
      levelSums[level] += root.val
      dfs(root.left, level + 1)
      dfs(root.right, level + 1)

    def replace(root: Optional[TreeNode], level: int, curr: Optional[TreeNode]) -> Optional[TreeNode]:
      nextLevel = level + 1
      nextLevelCousinsSum = (levelSums[nextLevel] if nextLevel < len(levelSums) else 0) - \
          (root.left.val if root.left else 0) - \
          (root.right.val if root.right else 0)
      if root.left:
        curr.left = TreeNode(nextLevelCousinsSum)
        replace(root.left, level + 1, curr.left)
      if root.right:
        curr.right = TreeNode(nextLevelCousinsSum)
        replace(root.right, level + 1, curr.right)
      return curr

    dfs(root, 0)
    return replace(root, 0, TreeNode(0))