2641. Cousins in Binary Tree II

Leetcode Question
Problem No.2641 Cousins in Binary Tree II
Problem Link- https://leetcode.com/problems/cousins-in-binary-tree-ii?envType=daily-question&envId=2024-10-23

Solution in C++

class Solution {
    public TreeNode replaceValueInTree(TreeNode root) {  
      updateVal(root, 0, levelSum(root));
      root.val = 0;
      return root;  
    }

    private void updateVal(TreeNode root, int level, List<Long> levelSum){
        if(root == null || level + 1 >= levelSum.size()) return;

//calculating cousins sum
long nextLevelSum = levelSum.get(level+1);
if(root.left!=null) nextLevelSum-=root.left.val;
if(root.right!=null) nextLevelSum-=root.right.val;

//update values of children
if(root.left!=null) root.left.val= (int)nextLevelSum;
if(root.right!=null) root.right.val= (int)nextLevelSum;

//call function updatVal for left and right child
updateVal(root.left, level+1, levelSum);
updateVal(root.right, level+1, levelSum);
    }

private List<Long> levelSum(TreeNode root){
    List<Long> ans = new ArrayList<>();
    Queue<TreeNode> q = new LinkedList<>();
    q.add(root);
    while(!q.isEmpty()) {
        int size = q.size();
        long sum = 0;
        for(int i=0; i<size; i++){
            TreeNode curr = q.remove();
            sum+=curr.val;
            if(curr.left!=null) q.add(curr.left);
            if(curr.right!=null) q.add(curr.right);
        }
        ans.add(sum);
    }
    return ans;
}    
}