Program to find k sublists with largest sums and return sums in ascending order in Python

Suppose we have a list of numbers called nums, and another value k, we have to find k sublists with the largest sums and return the sums in non-decreasing order.

So, if the input is like nums = [2, 4, 5, -100, 12, -30, 6, -2, 6] k = 3, then the output will be [10, 11, 12], as we have these 3 sublists with the largest sums ? [6, -2, 6], [2, 4, 5], [12].

Algorithm

To solve this, we will follow these steps ?

• Create a prefix sum array to calculate sublist sums efficiently
• Generate all possible sublist sums using prefix sums
• Use a max heap to keep track of the largest k sums
• Return the k largest sums in ascending order

Using Prefix Sum and Heap

Here's the implementation using prefix sums and a heap data structure ?

from heapq import heappop, heappush

class Solution:
    def solve(self, nums, k):
        # Create prefix sum array
        ps = [0 for _ in range(len(nums) + 1)]
        for i, v in enumerate(nums):
            ps[i + 1] = v + ps[i]
        
        # Use max heap (negative values for max heap behavior)
        hp = []
        for i in range(len(ps)):
            for j in range(i + 1, len(ps)):
                heappush(hp, -(ps[j] - ps[i]))
        
        # Extract k largest sums and return in ascending order
        return list(reversed([-heappop(hp) for _ in range(k)]))

# Test the solution
ob = Solution()
nums = [2, 4, 5, -100, 12, -30, 6, -2, 6]
k = 3
print(ob.solve(nums, k))

[10, 11, 12]

How It Works

The algorithm works in three main steps:

• Prefix Sum Calculation: We build a prefix sum array where ps[i] represents the sum of elements from index 0 to i-1. This allows us to calculate any sublist sum as ps[j] - ps[i].
• Heap Operations: We generate all possible sublist sums and store them as negative values in a min heap (to simulate max heap behavior).
• Result Extraction: We pop the k largest values from the heap and reverse the list to get ascending order.

Alternative Approach with Optimization

For better space efficiency, we can use a min heap of size k to track only the k largest sums ?

from heapq import heappop, heappush, heappushpop

def find_k_largest_sublists(nums, k):
    n = len(nums)
    # Create prefix sum array
    prefix_sum = [0] * (n + 1)
    for i in range(n):
        prefix_sum[i + 1] = prefix_sum[i] + nums[i]
    
    # Use min heap to maintain k largest sums
    min_heap = []
    
    for i in range(n + 1):
        for j in range(i + 1, n + 1):
            sublist_sum = prefix_sum[j] - prefix_sum[i]
            
            if len(min_heap) < k:
                heappush(min_heap, sublist_sum)
            elif sublist_sum > min_heap[0]:
                heappushpop(min_heap, sublist_sum)
    
    return sorted(min_heap)

# Test the optimized solution
nums = [2, 4, 5, -100, 12, -30, 6, -2, 6]
k = 3
result = find_k_largest_sublists(nums, k)
print(result)

[10, 11, 12]

Comparison

Conclusion

Use prefix sums to efficiently calculate sublist sums. The optimized min heap approach is more space-efficient for large arrays when k is small. Both methods return the k largest sublist sums in ascending order.