stacks_pick_m.py

# https://www.careercup.com/question?id=5711620404674560
# From N stacks pick M numbers.
# Since we are picking from stacks, any number popped MUST also be taken.
# Maximise the sum of numbers of popped.


def stacks_pick_m(stacks, M):
    dp = [[0 for _ in xrange(M + 1)] for _ in xrange(len(stacks) + 1)]

    for n, stack in enumerate(stacks, 1):
        for m in xrange(1, M + 1):
            dp[n][m] = max([dp[n-1][m-k] + sum(stack[0:k]) for k in xrange(m + 1)])

    # print DP table
    for row in dp:
        print row

stacks_pick_m([
    [5, 1, 2],
    [1, 11, 4, 6],
    [3, -2, 1, 4],
    [6, 2],
], 5)