Kth largest sum continuous subarrays

The problem, simply

You get an array of numbers and a number k. A subarray is an unbroken stretch of the array, like arr[1..3]. Every subarray has a sum. An array of n numbers has n (n + 1) / 2 subarrays. Find the k-th largest of all those sums.

For [3, 2, 1] the sums are 3, 5, 6, 2, 3 and 1. From big to small: 6, 5, 3, 3, 2, 1. With k = 2 the answer is 5.

A story to remember it

Picture it

Examples

Way 1: the simple way

Write every sum on the board, sort the board, and read the k-th from the top.

1. For each start i, keep a running sum while you move the end j to the right. Each step gives one subarray sum.
2. Put all sums in a list and sort it from big to small.
3. Return the item at index k - 1.

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

class Solution {
    public long kthLargestSum(int[] arr, int k) {
        List<Long> sums = new ArrayList<>();
        for (int i = 0; i < arr.length; i++) {
            long sum = 0;
            for (int j = i; j < arr.length; j++) {
                sum += arr[j];                 // the sum of arr[i..j]
                sums.add(sum);
            }
        }
        Collections.sort(sums, Collections.reverseOrder());
        return sums.get(k - 1);
    }
}

How long it takes: there are about n^2 / 2 sums, and sorting them costs O(n^2 log n). Extra room: the list of all sums, O(n^2). For n = 10,000 that is 50 million numbers.

Way 2: the clever way

Do not keep every sum. Keep only the k biggest so far in a min heap, so the smallest of them is on top.

1. Make the sums the same way, with a running sum for each start.
2. For each sum s:
   • If the heap holds fewer than k sums, add s.
   • Else if s is bigger than the top, remove the top and add s.
   • Else skip s.
3. The top of the heap is the k-th largest sum.

import java.util.PriorityQueue;

class Solution {
    public long kthLargestSum(int[] arr, int k) {
        PriorityQueue<Long> tray = new PriorityQueue<>();   // min heap of size k
        for (int i = 0; i < arr.length; i++) {
            long sum = 0;
            for (int j = i; j < arr.length; j++) {
                sum += arr[j];                              // sum of arr[i..j]
                if (tray.size() < k) {
                    tray.offer(sum);
                } else if (sum > tray.peek()) {
                    tray.poll();                            // smallest of the k leaves
                    tray.offer(sum);
                }
            }
        }
        return tray.peek();
    }
}

How long it takes: there are about n^2 / 2 sums and each costs at most O(log k) on the heap, so O(n^2 log k). Extra room: the heap, O(k).

Watch out for

• Use long for sums. Many big numbers added up can pass the int limit.
• A running sum makes each new sum O(1). Re-adding the whole stretch every time turns the work into O(n^3).
• k can be as big as n (n + 1) / 2. Then the heap holds every sum, and you are back to O(n^2) room.
• With negative numbers, sums can go down as the stretch grows. Do not stop early when the sum drops.

What interviewers ask next

• Can you beat O(n^2)? Yes, binary search on the answer. For a guess x, count the stretches with sum at least x using prefix sums and a merge sort count, O(n log n) per guess. The whole thing is O(n log n log R), where R is the range of sums.
• If all numbers are positive? Sums grow as the stretch grows, so for a guess x two pointers count stretches with sum at least x in O(n).
• The k-th smallest sum? Use a max heap of size k and skip sums that are bigger than the top.
• The k largest sums from two sorted arrays (pair sums)? Start a max heap with the best pair and push its two neighbors each time you pop, with a set to avoid repeats.

Remember this

Grown-up words: a run of wagons is a contiguous subarray, the running count is a prefix sum, the tray is a min heap (priority queue) of size k, and the front card is the heap top.