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].

To solve this, we will follow these steps −

• ps := a list of 1 + size of nums and fill with 0
• for each index i and value v nums, do
  • ps[i + 1] := v + ps[i]
• hp := a new list
• for i in range 0 to size of ps, do
  • for j in range i + 1 to size of ps, do
    • insert -(ps[j] - ps[i]) into ps heap
• res := pop all elements in ps heap and reverse them
• return res

Let us see the following implementation to get better understanding −

Example

 Live Demo

from heapq import heappop, heappush
class Solution:
   def solve(self, nums, k):
      ps = [0 for _ in range(len(nums) + 1)]
      for i, v in enumerate(nums):
         ps[i + 1] = v + ps[i]
      hp = []
      for i in range(len(ps)):
         for j in range(i + 1, len(ps)):
            heappush(hp, -(ps[j] - ps[i]))
            return list(reversed([-heappop(hp)
      for _ in range(k)]))
ob = Solution()
nums = [2, 4, 5, -100, 12, -30, 6, -2, 6] k = 3
print(ob.solve(nums, k))

Input

[2, 4, 5, -100, 12, -30, 6, -2, 6],3

Output

[10, 11, 12]