Count number of ways to construct a rhyming scheme

rhymes.py

# Jennivine's question about the number of rhyming schemes (AABA, etc.)
# for N lines using K letters.

# Given a function, return a cached version of that function.
def cache_function(fn):
    cache = dict()

    def fn_memo(*args):
        if args in cache:
            return cache[args]
        x = fn(*args)
        cache[args] = x
        return x

    return fn_memo


# f(n,k) is the number of rhyming schemes using exactly k letters.
#   f(n,1) = 1
#   f(n,k) = 0 if n < k
#   f(n,k) = 1 if n == k
#   f(n,k) = k(n-1,k-1) + k . f(n-1,k)
@cache_function
def f(n, k):
    if n < 1:  return 0
    if k == 1: return 1
    if n < k:  return 0
    if n == k: return 1
    return f(n-1,k-1) + k * f(n-1,k)

# Note to the reader: the two columns of code below are equivalent
#
#    @cache_function                 |
#    def f(x):                       |   def f(x):
#      time.sleep(10)                |     time.sleep(10)
#      return randint(x, x**2)       |     return randint(x, x**2)
#                                    |
#                                    |   f = cache_function(f)
#
# The @cache_function syntax is known as a "decorator". We "decorate" the function f by
# wrapping it inside another function. This is a convenient way to achieve caching.
# It's also good for logging/debugging, or checking/converting arguments before calling
# the function.

# The number of rhyming schemes for n lines using k letters is interpreted as
# "at most" k letters, so I need to add them up for all K in 1..k.
def nrhymes(n, k):
    return sum(f(n,i) for i in range(1,k+1))

for n,k in [(5,3), (5,5), (7,4), (10,2), (10,8), (50,23)]:
    print("N = %2d  K = %2d    #rhymes = %10d" % (n, k, nrhymes(n,k)))