LeetCode Entry

2583. Kth Largest Sum in a Binary Tree

22.10.2024 medium 2024 kotlin rust

kth largest level-sum in a tree

2583. Kth Largest Sum in a Binary Tree medium blog post substack youtube deep-dive

Join me on Telegram

https://t.me/leetcode_daily_unstoppable/776

Problem TLDR

kth largest level-sum in a tree #bfs #heap #quickselect

Intuition

To collect level sums we can use an iterative Breadth-First Search or a recursive Depth-First Search with level tracking.

To find kth largest, we can use a min-heap and maintain at most k items in it, or we can collect all the sums and then do a Quickselect algorithm to find kth largest value in O(n)

Approach

• it is simpler to store a non-null values in the queue
• in Rust we can destroy the tree with take or do a cheap Rc::clone (a simple .clone() call will do the recursive cloning and is slow)
• in c++ has built-in nth_element for Quickselect

Complexity

• Time complexity: \(O(n + log(n)log(k))\) or O(n) for Quickselect

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

Code

    fun kthLargestLevelSum(root: TreeNode?, k: Int): Long {
        val pq = PriorityQueue<Long>()
        val q = ArrayDeque<TreeNode>(listOf(root ?: return -1))
        while (q.size > 0) {
            pq += (1..q.size).sumOf { q.removeFirst().run {
                left?.let { q += it }; right?.let { q += it }; `val`.toLong() }}
            if (pq.size > k) pq.poll()
        }
        return if (pq.size == k) pq.poll() else -1
    }

    pub fn kth_largest_level_sum(root: Option<Rc<RefCell<TreeNode>>>, k: i32) -> i64 {
        let Some(r) = root else { return -1i64 };
        let (mut q, mut bh) = (VecDeque::from([r]), BinaryHeap::new());
        while q.len() > 0 {
            let sum = (0..q.len()).map(|_|{
                let n = q.pop_front().unwrap(); let n = n.borrow();
                if let Some(l) = &n.left { q.push_back(Rc::clone(l)) };
                if let Some(r) = &n.right { q.push_back(Rc::clone(r)) };
                n.val as i64
            }).sum::<i64>();
            bh.push(-sum); if bh.len() > k as usize { bh.pop(); }
        }
        if bh.len() == k as usize { -bh.pop().unwrap() } else { -1 }
    }

    long long kthLargestLevelSum(TreeNode* root, int k) {
        queue<TreeNode*>q; q.push(root); vector<long long> s;
        while (!q.empty()) {
            long long sum = 0;
            for (int i = q.size(); i; --i) {
                TreeNode* node = q.front(); q.pop(); sum += node->val;
                if (node->left) q.push(node->left);
                if (node->right) q.push(node->right);
            }
            s.push_back(sum);
        }
        return s.size() < k ? -1 : (nth_element(begin(s), begin(s) + k - 1, end(s), greater<>()), s[k-1]);
    }