castles.py

#!/usr/bin/env python3
#
# Answer to the 538 Riddler for Sept 9, 2022
#
# https://fivethirtyeight.com/features/can-you-cover-the-baking-sheet-with-cookies/

import sys


def find_combos(n_soldiers, n_castles):
    """
    Given a number of soldiers and number of castles, finds the number of ways you can
    distribute those soldiers across the castles.
    
    NB: Memoized to get it to run in any reasonable kind of time
    """
    if n_soldiers == 0 or n_castles < 2:
        return 1
    key = (n_soldiers, n_castles)
    if key not in find_combos.CACHE:
        # We can leave 0 to n_soldiers at this castle
        # So try each of those, adding up the number of possible combinations
        # tha each one allows
        total = 0
        for t in range(n_soldiers + 1):
            total += find_combos(n_soldiers - t, n_castles - 1)

        # Store it for later use
        find_combos.CACHE[key] = total
    return find_combos.CACHE[key]


find_combos.CACHE = {}
soldiers = int(sys.argv[1])
castles = int(sys.argv[2])

print(f"For {soldiers} soldiers over {castles} castles, there are {find_combos(soldiers, castles)} distributions.")