positive integer expansion for https://www.youtube.com/watch?v=PoSB9HK4kxg&ab_channel=MichaelPenn

positive-integer-expansion.py

# impl/test of https://www.youtube.com/watch?v=PoSB9HK4kxg&ab_channel=MichaelPenn
# m = \sum_{n=0}^{m-1} \sum_{k=0}^{n+1} \frac{(-1)^{k+1}}{n+1} \binom{n+1}{k} (k+1)^m
# this is from m=lim_{x->1} mx^{m-1} with the equivalent definition of
# derivative: f'(x)=ln(1+\Delta x)f(x)

from math import comb
from fractions import Fraction
from typing import Generator, Protocol


class Stringable(Protocol):
    def __str__(self):
        return self.__repr__()


def gen_pascal_triangle(line_count: int) -> list[list[int]]:
    lines = [[1]]
    for _ in range(line_count - 1):
        lines.append([1] + [a + b for a, b in zip(lines[-1], lines[-1][1:])] + [1])
    return lines


def gen_expansion(num: int) -> Generator[list[int], None, None]:
    yield from (
        (-1) ** (k + 1) * comb(n + 1, k) * (k + 1) ** num / (n + 1)
        for n in range(num)
        for k in range(n + 2)
    )


def gen_expansion_triangle(num: int) -> list[list[Fraction]]:
    # m = \sum_{n=0}^{m-1} \sum_{k=0}^{n+1} \frac{(-1)^{k+1}}{n+1} \binom{n+1}{k} (k+1)^m
    # we can regard the above formula as a triangle of n lines
    # the element at line `n` and index `k` is $\frac{(-1)^{k+1}}{n+1} \binom{n+1}{k} (k+1)^m$
    lines: list[list[Fraction]] = []
    for n in range(num):
        lines.append([])
        for k in range(n + 2):
            lines[-1].append(
                Fraction((-1) ** (k + 1) * comb(n + 1, k) * (k + 1) ** num, n + 1)
            )
    return lines


def print_triangle(triangle: list[list[Stringable]], print_sum: bool = False):
    # print a centered triangle of a stringable `triangle`
    # stringified elements in each line are separated by spaces
    # lines are centered
    rendered_elements: list[list[str]] = []
    max_ele_len = 0
    for line in triangle:
        rendered_elements.append([])
        for ele in line:
            rendered = str(ele)
            rendered_elements[-1].append(rendered)
            max_ele_len = max(max_ele_len, len(rendered))

    rendered_lines: list[str] = []
    for line in rendered_elements:
        rendered_lines.append("")
        for e in line:
            rendered_lines[-1] += f" {e:^{max_ele_len}} "

    max_line_len = len(rendered_lines[-1])
    total_sum = 0
    for line, raw_line in zip(rendered_lines, triangle):
        print(line.center(max_line_len) + f" |sum {sum(raw_line)}" * print_sum)
        if print_sum:
            total_sum += sum(raw_line)
    if print_sum:
        print(f"total sum: {total_sum}")


def verify_expansion(num: int):
    pt = gen_expansion_triangle(num)
    if num != sum(sum(line) for line in pt):
        raise ValueError("sum of expansion triangle is not equal to n")
    print_triangle(pt, True)


if __name__ == "__main__":
    verify_expansion(7)