Kth Largest Sum in a Binary Tree

Sanjeev SharmaSanjeev Sharma
1 min read

Advertisement

Problem

Return the kth largest sum among all level sums in a binary tree.

Solutions

# Python
from collections import deque
import heapq

def kthLargestLevelSum(root, k):
    q = deque([root])
    sums = []
    while q:
        level_sum = 0
        for _ in range(len(q)):
            node = q.popleft()
            level_sum += node.val
            if node.left:  q.append(node.left)
            if node.right: q.append(node.right)
        sums.append(level_sum)
    if k > len(sums):
        return -1
    return sorted(sums, reverse=True)[k-1]
// C++
long long kthLargestLevelSum(TreeNode* root, int k) {
    queue<TreeNode*>q; q.push(root);
    vector<long long>sums;
    while(!q.empty()){int sz=q.size();long long s=0;
        while(sz--){auto n=q.front();q.pop();s+=n->val;if(n->left)q.push(n->left);if(n->right)q.push(n->right);}
        sums.push_back(s);}
    if(k>(int)sums.size())return -1;
    nth_element(sums.begin(),sums.begin()+k-1,sums.end(),greater<long long>());
    return sums[k-1];
}
// Java
public long kthLargestLevelSum(TreeNode root, int k) {
    Queue<TreeNode>q=new LinkedList<>();q.add(root);
    PriorityQueue<Long>minH=new PriorityQueue<>();
    while(!q.isEmpty()){int sz=q.size();long s=0;
        while(sz-->0){TreeNode n=q.poll();s+=n.val;if(n.left!=null)q.add(n.left);if(n.right!=null)q.add(n.right);}
        minH.offer(s);if(minH.size()>k)minH.poll();}
    return minH.size()<k?-1:minH.peek();
}
// JavaScript
function kthLargestLevelSum(root, k) {
    let queue = [root], sums = [];
    while (queue.length) {
        sums.push(queue.reduce((s,n)=>s+n.val,0));
        queue = queue.flatMap(n=>[n.left,n.right].filter(Boolean));
    }
    if (k > sums.length) return -1;
    return sums.sort((a,b)=>b-a)[k-1];
}

Complexity

  • Time: O(n + L log L) where L = number of levels

Advertisement

Sanjeev Sharma

Written by

Sanjeev Sharma

Full Stack Engineer · E-mopro