HeapProblem 8 of 18

Kth largest sum continuous subarrays

Brute Force
Time: O(n² log n)
Space: O(n²)
Optimal
Time: O(n² log n)
Space: O(n)

Problem Statement

Given an array, find the Kth largest sum among all contiguous subarrays.

Example:

  • Input: arr = [3, 2, 1], k = 2
  • Output: 5
  • Explanation:
    • All subarray sums: [3], [2], [1], [3,2]=5, [2,1]=3, [3,2,1]=6
    • Sorted: [6, 5, 3, 3, 2, 1]
    • 2nd largest = 5

Approach 1: Generate All Sums and Sort (Brute Force)

Intuition

Generate all possible subarray sums, store them, sort in descending order, and return the Kth element.

Algorithm

  1. Generate all subarray sums - O(n²)
  2. Store all sums in an array
  3. Sort in descending order
  4. Return the Kth element
java
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

public class Solution {
    public int kthLargestSum(int[] arr, int k) {
        int n = arr.length;
        List<Integer> sums = new ArrayList<>();
        
        // Generate all subarray sums
        for (int i = 0; i < n; i++) {
            int sum = 0;
            for (int j = i; j < n; j++) {
                sum += arr[j];
                sums.add(sum);
            }
        }
        
        // Sort in descending order
        Collections.sort(sums, Collections.reverseOrder());
        
        // Return kth largest
        return sums.get(k - 1);
    }
}

Complexity Analysis

  • Time Complexity: O(n² log n²) = O(n² log n) - n² sums, sorting takes O(n² log n²)
  • Space Complexity: O(n²) - Storing all subarray sums

Approach 2: Using Min-Heap of Size K (Optimal)

Intuition

Instead of storing all sums, maintain a min-heap of size K. After processing all sums, the top of heap is the Kth largest.

Algorithm

  1. Generate subarray sums one by one
  2. Maintain a min-heap of size K
  3. For each sum:
    • If heap size < K, push sum
    • Else if sum > heap top, pop and push sum
  4. Heap top is the Kth largest
java
import java.util.PriorityQueue;

public class Solution {
    public int kthLargestSum(int[] arr, int k) {
        int n = arr.length;
        
        // Min-heap of size k
        PriorityQueue<Integer> minHeap = new PriorityQueue<>();
        
        // Generate all subarray sums
        for (int i = 0; i < n; i++) {
            int sum = 0;
            for (int j = i; j < n; j++) {
                sum += arr[j];
                
                if (minHeap.size() < k) {
                    minHeap.offer(sum);
                } else if (sum > minHeap.peek()) {
                    minHeap.poll();
                    minHeap.offer(sum);
                }
            }
        }
        
        return minHeap.peek();
    }
}

Complexity Analysis

  • Time Complexity: O(n² log k) - n² sums, each heap operation is O(log k)
  • Space Complexity: O(k) - Heap stores at most k elements

Approach 3: Using Max-Heap with Prefix Sums

Intuition

Use prefix sums to efficiently calculate subarray sums. Use a max-heap to extract K largest sums.

Algorithm

  1. Compute prefix sum array
  2. Generate all subarray sums using prefix sums: sum(i,j) = prefix[j+1] - prefix[i]
  3. Use max-heap and extract K elements
java
import java.util.PriorityQueue;
import java.util.Collections;

public class Solution {
    public long kthLargestSum(int[] arr, int k) {
        int n = arr.length;
        
        // Compute prefix sums
        long[] prefix = new long[n + 1];
        for (int i = 0; i < n; i++) {
            prefix[i + 1] = prefix[i] + arr[i];
        }
        
        // Max-heap
        PriorityQueue<Long> maxHeap = new PriorityQueue<>(Collections.reverseOrder());
        
        // Generate all subarray sums
        for (int i = 0; i < n; i++) {
            for (int j = i; j < n; j++) {
                long sum = prefix[j + 1] - prefix[i];
                maxHeap.offer(sum);
            }
        }
        
        // Extract k-1 elements
        for (int i = 0; i < k - 1; i++) {
            maxHeap.poll();
        }
        
        return maxHeap.peek();
    }
}

Complexity Analysis

  • Time Complexity: O(n² + k log n²) = O(n² log n) for large k
  • Space Complexity: O(n²) - Heap contains all sums

Approach 4: Binary Search + Counting (Advanced)

Intuition

Binary search on the answer. For a given value X, count how many subarray sums are >= X.

Algorithm

  1. Find range of possible sums [minSum, maxSum]
  2. Binary search for the Kth largest value
  3. Count function uses prefix sums and two pointers
java
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

public class Solution {
    private int countSubarrays(long[] prefix, long target) {
        int n = prefix.length - 1;
        int count = 0;
        
        for (int i = 0; i < n; i++) {
            for (int j = i; j < n; j++) {
                if (prefix[j + 1] - prefix[i] >= target) {
                    count++;
                }
            }
        }
        
        return count;
    }
    
    public long kthLargestSum(int[] arr, int k) {
        int n = arr.length;
        
        // Compute prefix sums
        long[] prefix = new long[n + 1];
        for (int i = 0; i < n; i++) {
            prefix[i + 1] = prefix[i] + arr[i];
        }
        
        // Get all sums
        List<Long> allSums = new ArrayList<>();
        for (int i = 0; i < n; i++) {
            for (int j = i; j < n; j++) {
                allSums.add(prefix[j + 1] - prefix[i]);
            }
        }
        
        Collections.sort(allSums, Collections.reverseOrder());
        return allSums.get(k - 1);
    }
}

Complexity Analysis

  • Time Complexity: O(n² log(maxSum - minSum)) for binary search approach
  • Space Complexity: O(n) for prefix array

Key Takeaways

  1. There are n(n+1)/2 subarrays, hence n² sums to consider
  2. Min-heap of size K reduces space from O(n²) to O(k)
  3. Prefix sums allow O(1) calculation of any subarray sum
  4. For small K, min-heap approach is efficient
  5. For large K close to n², max-heap with extraction might be better
  6. Binary search approach exists but doesn't improve complexity much for this problem
  7. This is a variation of "Kth largest element" applied to subarray sums