roots.py

import math
from functools import cache
from typing import Optional


@cache
def is_prime(n: int, primes: tuple[int]) -> bool:
    sqrt = math.floor(math.sqrt(n))

    for i in primes:
        if i > sqrt:
            break
        if (n % i) == 0:
            return False

    return True


@cache
def primes_below(n: int) -> list[int]:
    primes = []
    for i in range(2, n):
        if is_prime(i, primes=tuple(primes)):
            primes.append(i)

    return primes


@cache
def lowest_num_to_look(n: int) -> int:
    return math.ceil(math.sqrt(n)) + 1


def find_roots_(
    n: int, roots: Optional[list[int]] = None, primes: Optional[list[int]] = None
) -> list[int]:
    if roots is None:
        roots = [1]
    if primes is None:
        primes = primes_below(lowest_num_to_look(n))

    var_n = n
    new_roots = []

    for prime in filter(lambda x: x < lowest_num_to_look(n), primes):
        if (var_n % prime) == 0:
            new_roots.append(prime)
            var_n //= prime

    if new_roots:
        roots.extend(new_roots)
        return find_roots_(var_n, roots, primes)
    else:
        if var_n != 1:
            roots.append(var_n)
        return sorted(roots)


def find_roots(n: int) -> list[int]:
    if not isinstance(n, int) or not (n >= 1):
        raise ValueError("Please enter an integer number >= 1.")

    return find_roots_(n)


if __name__ == "__main__":
    print(find_roots(987654321))

    # from functools import reduce
    #
    # for i in range(1, 1_000_001):
    #     if reduce(lambda x, y: x * y, (ans := find_roots(i))) != i:
    #         print(f"Error! for {i:,} function returned: {ans}")