Dynamic Programming: Taking Coins

coins.py

# DP practice problem.
# Feb 21, 2017.
# Jason Sun sunapi386.ca
#
# Problem
#
# You are a given collection of coins, where each stack of coins is of h height, and there are s stacks.
# Take exactly h coins, such that the sum of the coins is the maximum possible.
# Caveat: if you take a coin from a stack, all the coins above it must be taken too.
#
# Example:
# s = 3
# h = 3
# coins = [[1,2,3], [4,5,6], [7,8,9]]
#
# The different ways of taking 3 coins are:
# 1,2,3 # entire first stack
# 1,2,4 # the top two from first stack, 2 options
# 1,2,7
# 1,4,5 # only the top one from first stack, 3 options
# 1,4,7
# 1,7,8
#
# 4,5,6 # entire second stack
# 4,5,1
# 4,5,7
# 4,1,7
# 4,1,2
# 4,7,8
#
# 7,8,9 # entire third stack
# 7,8,1
# 7,8,4
# 7,1,2
# 7,1,4
# 7,4,5

# Assume coins are in denominations of 1 to 10 value.
from random import randint
def coin():
  return randint(1, 10)

s = 4
h = 3

# Build s stacks of h coins (represented as a 2D array).
coins = [ [coin() for _ in range(h) ] for _ in range(s) ]
# coins = [[5, 8, 8], [6, 3, 1], [3, 6, 3], [8, 5, 1]]

#
# Solution
#

# Build a sums table, to represent how much value we get from taking the next coin.
# This is to help us do calculations.
from itertools import accumulate, chain
sums = [list(accumulate(stack)) for stack in coins]
# sums = [[5, 13, 21], [6, 9, 10], [3, 9, 12], [8, 13, 14]]

def reverseOneZipCoins(A,B,n):
  return list(zip([0]+A[:n],reversed([0]+B[:n])))

def maxSum(A):
  return max([sum(tuple) for tuple in A])

def merge(A, B):
  return [maxSum(reverseOneZipCoins(A,B,n)) for n in range(1,len(A)+1)]


# We can take the sums and merge them together, building a DP table.
# Each of row represents the optimal solution for the first i coin stack.
# So table[-1] represents having looked at all of the stack.
#
table = list(accumulate(sums, lambda acc, cur: merge(acc,cur)))
# table = [[5, 13, 21], [6, 13, 21], [6, 13, 21], [8, 14, 21]]

# For taking:
#   - 1 coin, max taken is 8.
#   - 2 coin, 14.
#   - 3 coin, 21.
optimal = table[-1]
# optimal = [8, 14, 21]

print("coins=", coins)
print("sums=", sums)
print("table=", table)
print("optimal=", optimal)