LeetCode Entry

2641. Cousins in Binary Tree II

23.10.2024 medium 2024 kotlin rust

Replace Tree's values with cousines sum

2641. Cousins in Binary Tree II medium blog post substack youtube deep-dive

Join me on Telegram

https://t.me/leetcode_daily_unstoppable/777

Problem TLDR

Replace Tree’s values with cousines sum #medium #bfs

Intuition

First, understand the problem, we only care about the current level’s row:

Now, the task is to traverse Tree level by level and precompute the total next level sum and the current parent's sum.

Approach

• consider only the current and the next level
• we can modify at the same time as adding to the queue

Complexity

• Time complexity: \(O(n)\)

• Space complexity: \(O(n)\)

Code

    fun replaceValueInTree(root: TreeNode?): TreeNode? {
        val q = ArrayDeque<TreeNode>(listOf(root ?: return root))
        while (q.size > 0) {
            val sum = q.sumBy { (it.left?.`val` ?: 0) + (it.right?.`val` ?: 0) }
            repeat(q.size) { q.removeFirst().run {
                var nv = sum - (left?.`val` ?: 0) - (right?.`val` ?: 0)
                left?.let { it.`val` = nv; q += it }
                right?.let { it.`val` = nv; q += it }
            }}
        }
        return root.also { it.`val` = 0 }
    }

    pub fn replace_value_in_tree(mut root: Option<Rc<RefCell<TreeNode>>>) -> Option<Rc<RefCell<TreeNode>>> {
        let Some(r) = root.clone() else { return root }; let mut q = VecDeque::from([r]);
        while q.len() > 0 {
            let mut sum = q.iter().map(|n| { let n = n.borrow();
                n.left.as_ref().map_or(0, |n| n.borrow().val) +
                n.right.as_ref().map_or(0, |n| n.borrow().val)}).sum::<i32>();
            for _ in 0..q.len() {
                let n = q.pop_front().unwrap(); let mut n = n.borrow_mut();
                let mut s =  sum - n.left.as_ref().map_or(0, |n| n.borrow().val) -
                             n.right.as_ref().map_or(0, |n| n.borrow().val);
                if let Some(l) = n.left.clone() { l.borrow_mut().val = s; q.push_back(l); }
                if let Some(r) = n.right.clone() { r.borrow_mut().val = s; q.push_back(r); }
            }
        }
        if let Some(r) = &root { r.borrow_mut().val = 0 }; root
    }

    TreeNode* replaceValueInTree(TreeNode* root) {
        if (!root) return root; queue<TreeNode*> q({root}); root->val = 0;
        while (!q.empty()) {
            int sum = 0, size = q.size();
            for (int i = 0; i < size; ++i, q.push(q.front()), q.pop()) {
                auto node = q.front();
                sum += (node->left ? node->left->val : 0) + (node->right ? node->right->val : 0);
            }
            for (int i = 0; i < size; ++i) {
                auto node = q.front(); q.pop();
                int nv = sum - (node->left ? node->left->val : 0) - (node->right ? node->right->val : 0);
                if (node->left) node->left->val = nv, q.push(node->left);
                if (node->right) node->right->val = nv, q.push(node->right);
            }
        }
        return root;
    }