Kth Largest Sum in a Binary Tree

medium tree bfs heap sorting

Problem

Given the root of a binary tree and an integer k, a level sum is the total of all node values at one depth of the tree. Return the k^th largest of these level sums. If the tree has fewer than k levels, return -1.

Inputtree = [5, 8, 9, 2, 1, 3, 7, 4, 6], k = 2

Output13

Level sums are 5, 17, 13, 10. Sorted high to low: 17, 13, 10, 5. The 2nd largest is 13.

tree (level-order)

def kth_largest_level_sum(root, k):
    sums = []
    level = [root]
    while level:
        total = sum(n.val for n in level)
        sums.append(total)
        nxt = []
        for n in level:
            if n.left: nxt.append(n.left)
            if n.right: nxt.append(n.right)
        level = nxt
    if len(sums) < k:
        return -1
    sums.sort(reverse=True)
    return sums[k - 1]

function kthLargestLevelSum(root, k) {
  const sums = [];
  let level = [root];
  while (level.length) {
    const total = level.reduce((s, n) => s + n.val, 0);
    sums.push(total);
    const next = [];
    for (const n of level) {
      if (n.left) next.push(n.left);
      if (n.right) next.push(n.right);
    }
    level = next;
  }
  if (sums.length < k) return -1;
  sums.sort((a, b) => b - a);
  return sums[k - 1];
}

class Solution {
    public long kthLargestLevelSum(TreeNode root, int k) {
        List<Long> sums = new ArrayList<>();
        List<TreeNode> level = new ArrayList<>(); level.add(root);
        while (!level.isEmpty()) {
            long total = 0; for (TreeNode n : level) total += n.val;
            sums.add(total);
            List<TreeNode> next = new ArrayList<>();
            for (TreeNode n : level) {
                if (n.left != null) next.add(n.left);
                if (n.right != null) next.add(n.right);
            }
            level = next;
        }
        if (sums.size() < k) return -1;
        sums.sort(Collections.reverseOrder());
        return sums.get(k - 1);
    }
}

long long kthLargestLevelSum(TreeNode* root, int k) {
    vector<long long> sums;
    vector<TreeNode*> level = { root };
    while (!level.empty()) {
        long long total = 0; for (auto* n : level) total += n->val;
        sums.push_back(total);
        vector<TreeNode*> next;
        for (auto* n : level) {
            if (n->left) next.push_back(n->left);
            if (n->right) next.push_back(n->right);
        }
        level = next;
    }
    if ((int)sums.size() < k) return -1;
    sort(sums.rbegin(), sums.rend());
    return sums[k - 1];
}

Explanation

The problem talks about whole levels of the tree, so the natural tool is breadth-first search (BFS), which walks the tree one depth at a time.

We hold the current depth's nodes in a list called level. We add up their values to get that level's sum and push it onto a list called sums. To move down, we build the next level from every current node's left and right children. The loop ends when a level has no nodes.

After the sweep, sums contains exactly one number per level. The k^th largest level sum is now just the k^th largest entry of that small list. Sorting it in descending order and reading position k - 1 gives the answer.

If the tree has fewer levels than k (so sums has fewer than k entries), there is no k^th largest sum, and we return -1.

Example: [5, 8, 9, 2, 1, 3, 7, 4, 6], k = 2. Level sums are 5, 17, 13, 10. Sorted high to low that is 17, 13, 10, 5, and the 2nd value is 13.

If the value count is large, a min-heap of size k can replace the sort, but for the handful of levels a tree usually has, a plain sort is simplest.

Time: O(n + L log L) Space: O(w)