Skip to content

Instantly share code, notes, and snippets.

@krone

krone/stacks_pick_m.py

Created November 20, 2017 07:49
Show Gist options
  • Save krone/89ecf2c2472e403e180564169b89eacf to your computer and use it in GitHub Desktop.
Save krone/89ecf2c2472e403e180564169b89eacf to your computer and use it in GitHub Desktop.
Download ZIP
Raw
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)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment