whole number solution for sqrt(a)+sqrt(b) = sqrt(n), using integer math only (youtube video pxHd8tLI65Q)

output

solving for 2023
((0, 2023), (7, 1792), (28, 1575), (63, 1372), (112, 1183), (175, 1008), (252, 847), (343, 700), (448, 567), (567, 448), (700, 343), (847, 252), (1008, 175), (1183, 112), (1372, 63), (1575, 28), (1792, 7), (2023, 0))
solving for 897289125379227
num results for the large number test: 17294404

Process finished with exit code 0

solution.py

def prime_factors(n):
    i = 2
    while i * i <= n:
        if n % i:
            i += 1
        else:
            n //= i
            yield i
    if n > 1:
        yield n


def sqrt(n):
    """
    returns two integer numbers (a,b), where a*sqrt(b)=sqrt(n)
    """
    factors = tuple(prime_factors(n))

    whole = 1

    # primes are already sorted; look for matching pairs (square roots) and extract them outside the root
    i = 0
    while i < len(factors) - 1:
        f1, f2 = factors[i:i + 2]
        if f1 == f2:
            whole *= f1
            i += 1  # skip this pair for next comparison
        i += 1

    return whole, n // (whole*whole)


def solve(n):
    print("solving for",n)
    a,b = sqrt(n)

    def solve(x):
        return b*x*x

    for i in range(a+1):
        yield solve(i),solve(a-i)

print(tuple(solve(2023)))


# solve for reasonably large number (897289125379227) just to check performance:) dont print results
print("num results for the large number test:", len(tuple(solve(7**13*21**3))))