Bagels Problem Solver

bagels.py

import argparse

# The approach I thought of first. It uses a DP technique
# where the number of ways to fit t types of bagels into
# s slots is computed for every t up to and inluding the
# number of bagel types and for every s up to and including
# the number of slots to fill.
# Running time: Theta(types*(slots)^2) (pseudo-polynomial).
def calculate_ways(types, slots):
    if types == 0: return 0
    if types == 1: return 1

    dp = [[0 for _ in range(slots+1)] for _ in range(types+1)]

    for s in range(0,slots+1):
        dp[1][s] = 1

    for t in range(2,types+1):
        for s in range(0, slots+1):
            dp[t][s] = sum(dp[t-1][s-x] for x in range(0,s+1))

    return dp[types][slots]

# This uses the same logic as the DP solution, but it does the
# recursion explicitly instead.
# Running time: Theta(answer(types, slots)) (roughly exponential in the minimum of the two arguments).
def recursive_version(types, slots):
    if types == 1: return 1
    return sum(recursive_version(types-1, slots-s) for s in range(0,slots+1))

# This computes the answer in a brute force kind of way by iterating through
# all of the different ways to partition the bagels.
# Running time: Theta(answer(types, slots)) (roughly exponential in the minimum of the two arguments).
def brute_force(types, slots):
    val = [0 for _ in range(slots)]
    ways = 1

    while True:
        i = 0

        while i < slots:
            val[i] += 1
            if val[i] < types:
                break
            else:
                i += 1

        if i == slots:
            break
        else:
            for x in range(0, i):
                val[x] = val[i]

        ways += 1

    return ways

# This uses the closed form solution discussed in the Numberphile video: https://youtu.be/UTCScjoPymA
# Running time: Theta(slots+types) (pseudo-polynomial).
def closed_form(types, slots):
     return factorial(slots+types-1)/(factorial(types-1)*factorial(slots))

# Factorial helper function.
# Running time: Theta(n) (pseudo-polynomial).
def factorial(n):
    return reduce(lambda acc, x: x*acc, range(1,n+1), 1)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('types', type=int, help='the number of bagel types')
    parser.add_argument('slots', type=int, help='the number of bagels we are buying')
    parser.add_argument('--dp', action='store_true', help='do it using the dp solution')
    parser.add_argument('--recur', action='store_true', help='do it recursively as well')
    parser.add_argument('--brute-force', action='store_true', help='do it using brute force')
    parser.add_argument('--closed-form', action='store_true', help='do it using the boring closed form solution')
    args = parser.parse_args()

    # Used the closed form to make this fast.
    print 'Ways to place', args.types, 'types of bagels into', args.slots, 'slots:', closed_form(args.types, args.slots)

    if args.dp:
        print 'DP answer:', calculate_ways(args.types, args.slots)
    if args.recur:
        print 'Recurive answer:', recursive_version(args.types, args.slots)
    if args.brute_force:
        print 'Brute force answer:', brute_force(args.types, args.slots)
    if args.closed_form:
        print 'Closed form answer:', closed_form(args.types, args.slots)