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
- for j in range i + 1 to size of ps, do
- res := pop all elements in ps heap and reverse them
- return res
Let us see the following implementation to get better understanding −
Example
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]