Program to find maximum sum of popped k elements from a list of stacks in Python

Suppose we have a list of stacks and an integer k. We have to find the maximum possible sum that can be achieved from popping off exactly k elements from any combination of the stacks.

So, if the input is like stacks = [[50, -4, -15],[2],[6, 7, 8]], k = 4, then the output will be 39, as we can pop off all 3 elements from the first stack and pop the last element of the last stack to get -15 + -4 + 50 + 8 = 39.

To solve this, we will follow these steps −

• Define a function rec() . This will take i, n

• if n is same as k, then

  • return 0

• if n > k, then

  • return negative infinity

• if i is same as count of stacks , then

  • return negative infinity

• if i is same as count of stacks - 1, then

  • needed := k - n

  • if needed > element count of stacks[i], then

    • return negative infinity

  • otherwise,

    • return sum of elements of stack[i] last needed number of elements

• res := -math.inf, su := 0

• for sti in range size of stacks[i] - 1 to 0,

  decrease by 1, do

  • su := su + stacks[i, sti]

  • localres := su + rec(i + 1, n + size of stacks[i] - sti)

  • res := maximum of res and localres

• return maximum of res and rec(i + 1, n)

• From the main method call rec(0, 0)

Let us see the following implementation to get better understanding −

Example

import math

class Solution:
   def solve(self, stacks, k):
      def rec(i, n):
         if n == k:
            return 0
         if n > k:
            return -math.inf
         if i == len(stacks):
            return -math.inf
         if i == len(stacks) - 1:
            needed = k - n
            if needed > len(stacks[i]):
               return -math.inf
            else:
               return sum(stacks[i][-needed:])
         res, su = -math.inf, 0
         for sti in range(len(stacks[i]) - 1, -1, -1):
            su += stacks[i][sti]
            localres = su + rec(i + 1, n + len(stacks[i]) - sti)
            res = max(res, localres)
         return max(res, rec(i + 1, n))

      return rec(0, 0)

ob = Solution()
stacks = [
   [50, -4, -15],
   [2],
   [6, 7, 8]
]
k = 4
print(ob.solve(stacks, k))

Input

[[50, -4, -15],[2],[6, 7, 8]], 4

Output

39