partition.py

def pentagonal(n):
    return n * (3 * n - 1) // 2


@lru_cache(maxsize=None)
def generalized_pentagonal(n):
    if n == 0:
        return pentagonal(0)
    if n % 2 == 0:
        return pentagonal(-n // 2)
    else:
        return pentagonal(n // 2 + 1)


@lru_cache(maxsize=None)
def partition(n):
    if n < 0:
        return 0
    if n == 0:
        return 1
    result = 0
    m = 0
    penta = 0
    while penta <= n:
        m += 1
        sign = -1 if (m - 1) % 4 > 1 else 1
        penta = generalized_pentagonal(m)
        current = sign * partition(n - penta)
        result += current
        if current == 0:
            break
    return result