How (and why) does "Python compute square roots of perfect squares exactly"?

Viewed 362

This amazing code golf answer to Is this number Loeschian? in its entirety:

Python, 49 bytes

lambda n:0in[(n-3*i*i+0j)**.5%1for i in range(n)]

Uses the equivalent quadratic form given on OEIS of n == 3*i*i+j*j. Check whether n-3*i*i is a perfect square for any i by taking its square root and checking if it's an integer, i.e. equals 0 modulo 1. Note that Python computes square roots of perfect squares exactly, without floating point error. The +0j makes it a complex number to avoid an error on the square root of a negative.

How does Python do this? Does **.5 "detect" that a given number is a perfect square somehow? Is this only reliable for integer input or will it work on floats up to some size as well?

I've also added a parenthetical Why? to the question; is this something that programmers rely upon? Is it for speed? Does it come with a cost?

1 Answers

You can check out the source code here. They describe the algorithm they use for computing the (approximate) square root of nonnegative integers, and show that for perfect squares the algorithm gives the exact answer. The code is C, but they give a translation of the code into Python:

def isqrt(n):
    """
    Return the integer part of the square root of the input.
    """
    n = operator.index(n)
    if n < 0:
        raise ValueError("isqrt() argument must be nonnegative")
    if n == 0:
        return 0
    c = (n.bit_length() - 1) // 2
    a = 1
    d = 0
    for s in reversed(range(c.bit_length())):
        # Loop invariant: (a-1)**2 < (n >> 2*(c - d)) < (a+1)**2
        e = d
        d = c >> s
        a = (a << d - e - 1) + (n >> 2*c - e - d + 1) // a
    return a - (a*a > n)

I assume, but haven't yet checked, that when computing a power at runtime, Python first checks that 1. the base is a nonnegative integer, 2. the exponent is exactly 0.5, and if those both hold then it invokes the code I linked to above.

Related