isqrt.py

def isqrt(n):
    """
    Compute the integer square root of an integer n. 
    This is equivalent to `floor(sqrt(n))` for small n, but is correct
    for integers of arbitrary size, unlike the floating point version.
    This implementation is a rewritten version of the one provided by
    nibot in this Stack Overflow answer:
        <http://stackoverflow.com/a/23279113/2738025>
    @author Alexander Gosselin
    @email  alexandergosselin@gmail.com
    @date   September 25, 2016
    """
    a = 0 # a is the current answer.
    r = 0 # r is the current remainder n - a**2.
    for s in reversed(range(0, n.bit_length(), 2)): # Shift n by s bits.
        t = n >> s & 3 # t is the two next most significant bits of n.
        r = r << 2 | t # Increase the remainder as if no new bit is set.
        c = a << 2 | 1 # c is an intermediate value used for comparison.
        b = r >= c     # b is the next bit in the remainder.
        if b: 
            r -= c     # b has been set, so reduce the remainder.
        a = a << 1 | b # Update the answer to include b.
    return (a, r)

def test_isqrt(trials=100, root_bits=128):
    import random
    for _ in range(trials):
        root = int(random.getrandbits(root_bits))
        assert root == isqrt(root**2)[0]