Predicting a number sequence from an unknown polynomial by finite differences

predict.py

from rich import print
from rich.table import Table


def subsequent_pairs(L):
    L = iter(L)
    a = next(L)
    b = next(L)
    while True:
        yield (a, b)
        a = b
        try:
            b = next(L)
        except StopIteration:
            return


def pretty_print(data, color_last=True):
    table = Table(expand=False)

    table.add_column(justify="right")
    chosen_len = len(data[0])-1 if color_last else len(data[0])
    for i in range(chosen_len):
        table.add_column(f"f({i})", justify="center")
    if color_last:
        table.add_column("prediction", justify="right")

    for i, row in enumerate(data):
        row = [f"Δ{i}"] + [str(x) for x in row]
        if color_last:
            row[-1] = f"[magenta]{row[-1]}[/magenta]"
        table.add_row(*row)
    print(table)


def predict(data, print_table=False):
    difference_table = [data[:]]
    while True:
        difference_table.append(
            [y - x for (x, y) in subsequent_pairs(difference_table[-1])]
        )

        if all(x == 0 for x in difference_table[-1]):
            break  # success

        if len(difference_table[-1]) <= 1:
            if print_table:
                pretty_print(difference_table, color_last=False)
                print("Either not a polynomial, or else not enough data.")
            return None

    difference_table[-1].append(0)
    for i in reversed(range(len(difference_table)-1)):
        difference_table[i].append(
            difference_table[i+1][-1] + difference_table[i][-1])

    if print_table:
        pretty_print(difference_table)

    return difference_table[0][-1]


if __name__ == "__main__":
    sequence = "1 8 27 64 125 216"
    print(sequence + " ?")
    sequence = [int(x) for x in sequence.split()]
    predict(sequence, print_table=True)