Problem statement

https://binarysearch.com/problems/K-Stack-Pops/

Solution

Let use dp with states dp(m, k), where it is (current index of stack, nubmer of elements we still need to take). Then on each state we can try to take 0, ..., min(k, len(stacks[m])) elements from stack m.

Complexity

It is O(N * K * M) for time and O(N * K) for space.

Code

class Solution:
    def solve(self, stacks, K):
        N = len(stacks)
        @lru_cache(None)
        def dp(n, k):
            if n == N:
                if k == 0: return 0
                if k > 0: return -float("inf")
            ans = dp(n + 1, k)
            sm = 0
            for i in range(min(k, len(stacks[n]))):
                sm += stacks[n][~i]
                ans = max(ans, dp(n + 1, k - i - 1) + sm)
            return ans

        return dp(0, K)