'Integer square root in python

Is there an integer square root somewhere in python, or in standard libraries? I want it to be exact (i.e. return an integer), and bark if there's no solution.

At the moment I rolled my own naive one:

def isqrt(n):
    i = int(math.sqrt(n) + 0.5)
    if i**2 == n:
        return i
    raise ValueError('input was not a perfect square')

But it's ugly and I don't really trust it for large integers. I could iterate through the squares and give up if I've exceeded the value, but I assume it would be kinda slow to do something like that. Also I guess I'd probably be reinventing the wheel, something like this must surely exist in python already...

pythonmathintegersqrt


Solution 1:[1]

Update: Python 3.8 has a math.isqrt function in the standard library!

I benchmarked every (correct) function here on both small (0…222) and large (250001) inputs. The clear winners in both cases are gmpy2.isqrt suggested by mathmandan in first place, followed by Python 3.8’s math.isqrt in second, followed by the ActiveState recipe linked by NPE in third. The ActiveState recipe has a bunch of divisions that can be replaced by shifts, which makes it a bit faster (but still behind the native functions):

def isqrt(n):
    if n > 0:
        x = 1 << (n.bit_length() + 1 >> 1)
        while True:
            y = (x + n // x) >> 1
            if y >= x:
                return x
            x = y
    elif n == 0:
        return 0
    else:
        raise ValueError("square root not defined for negative numbers")

Benchmark results:

(* Since gmpy2.isqrt returns a gmpy2.mpz object, which behaves mostly but not exactly like an int, you may need to convert it back to an int for some uses.)

Solution 2:[2]

Sorry for the very late response; I just stumbled onto this page. In case anyone visits this page in the future, the python module gmpy2 is designed to work with very large inputs, and includes among other things an integer square root function.

Example:

>>> import gmpy2
>>> gmpy2.isqrt((10**100+1)**2)
mpz(10000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000001L)
>>> gmpy2.isqrt((10**100+1)**2 - 1)
mpz(10000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000L)

Granted, everything will have the "mpz" tag, but mpz's are compatible with int's:

>>> gmpy2.mpz(3)*4
mpz(12)

>>> int(gmpy2.mpz(12))
12

See my other answer for a discussion of this method's performance relative to some other answers to this question.

Download: https://code.google.com/p/gmpy/

Solution 3:[3]

Long-hand square root algorithm

It turns out that there is an algorithm for computing square roots that you can compute by hand, something like long-division. Each iteration of the algorithm produces exactly one digit of the resulting square root while consuming two digits of the number whose square root you seek. While the "long hand" version of the algorithm is specified in decimal, it works in any base, with binary being simplest to implement and perhaps the fastest to execute (depending on the underlying bignum representation).

Because this algorithm operates on numbers digit-by-digit, it produces exact results for arbitrarily large perfect squares, and for non-perfect-squares, can produce as many digits of precision (to the right of the decimal place) as desired.

There are two nice writeups on the "Dr. Math" site that explain the algorithm:

And here's an implementation in Python:

def exact_sqrt(x):
    """Calculate the square root of an arbitrarily large integer. 

    The result of exact_sqrt(x) is a tuple (a, r) such that a**2 + r = x, where
    a is the largest integer such that a**2 <= x, and r is the "remainder".  If
    x is a perfect square, then r will be zero.

    The algorithm used is the "long-hand square root" algorithm, as described at
    http://mathforum.org/library/drmath/view/52656.html

    Tobin Fricke 2014-04-23
    Max Planck Institute for Gravitational Physics
    Hannover, Germany
    """

    N = 0   # Problem so far
    a = 0   # Solution so far

    # We'll process the number two bits at a time, starting at the MSB
    L = x.bit_length()
    L += (L % 2)          # Round up to the next even number

    for i in xrange(L, -1, -1):

        # Get the next group of two bits
        n = (x >> (2*i)) & 0b11

        # Check whether we can reduce the remainder
        if ((N - a*a) << 2) + n >= (a<<2) + 1:
            b = 1
        else:
            b = 0

        a = (a << 1) | b   # Concatenate the next bit of the solution
        N = (N << 2) | n   # Concatenate the next bit of the problem

    return (a, N-a*a)

You could easily modify this function to conduct additional iterations to calculate the fractional part of the square root. I was most interested in computing roots of large perfect squares.

I'm not sure how this compares to the "integer Newton's method" algorithm. I suspect that Newton's method is faster, since it can in principle generate multiple bits of the solution in one iteration, while the "long hand" algorithm generates exactly one bit of the solution per iteration.

Source repo: https://gist.github.com/tobin/11233492

Solution 4:[4]

Here's a very straightforward implementation:

def i_sqrt(n):
    i = n.bit_length() >> 1    # i = floor( (1 + floor(log_2(n))) / 2 )
    m = 1 << i    # m = 2^i
    #
    # Fact: (2^(i + 1))^2 > n, so m has at least as many bits 
    # as the floor of the square root of n.
    #
    # Proof: (2^(i+1))^2 = 2^(2i + 2) >= 2^(floor(log_2(n)) + 2)
    # >= 2^(ceil(log_2(n) + 1) >= 2^(log_2(n) + 1) > 2^(log_2(n)) = n. QED.
    #
    while m*m > n:
        m >>= 1
        i -= 1
    for k in xrange(i-1, -1, -1):
        x = m | (1 << k)
        if x*x <= n:
            m = x
    return m

This is just a binary search. Initialize the value m to be the largest power of 2 that does not exceed the square root, then check whether each smaller bit can be set while keeping the result no larger than the square root. (Check the bits one at a time, in descending order.)

For reasonably large values of n (say, around 10**6000, or around 20000 bits), this seems to be:

All of these approaches succeed on inputs of this size, but on my machine, this function takes around 1.5 seconds, while @Nibot's takes about 0.9 seconds, @user448810's takes around 19 seconds, and the gmpy2 built-in method takes less than a millisecond(!). Example:

>>> import random
>>> import timeit
>>> import gmpy2
>>> r = random.getrandbits
>>> t = timeit.timeit
>>> t('i_sqrt(r(20000))', 'from __main__ import *', number = 5)/5. # This function
1.5102493192883117
>>> t('exact_sqrt(r(20000))', 'from __main__ import *', number = 5)/5. # Nibot
0.8952787937686366
>>> t('isqrt(r(20000))', 'from __main__ import *', number = 5)/5. # user448810
19.326695976676184
>>> t('gmpy2.isqrt(r(20000))', 'from __main__ import *', number = 5)/5. # gmpy2
0.0003599147067689046
>>> all(i_sqrt(n)==isqrt(n)==exact_sqrt(n)[0]==int(gmpy2.isqrt(n)) for n in (r(1500) for i in xrange(1500)))
True

This function can be generalized easily, though it's not quite as nice because I don't have quite as precise of an initial guess for m:

def i_root(num, root, report_exactness = True):
    i = num.bit_length() / root
    m = 1 << i
    while m ** root < num:
        m <<= 1
        i += 1
    while m ** root > num:
        m >>= 1
        i -= 1
    for k in xrange(i-1, -1, -1):
        x = m | (1 << k)
        if x ** root <= num:
            m = x
    if report_exactness:
        return m, m ** root == num
    return m

However, note that gmpy2 also has an i_root method.

In fact this method could be adapted and applied to any (nonnegative, increasing) function f to determine an "integer inverse of f". However, to choose an efficient initial value of m you'd still want to know something about f.

Edit: Thanks to @Greggo for pointing out that the i_sqrt function can be rewritten to avoid using any multiplications. This yields an impressive performance boost!

def improved_i_sqrt(n):
    assert n >= 0
    if n == 0:
        return 0
    i = n.bit_length() >> 1    # i = floor( (1 + floor(log_2(n))) / 2 )
    m = 1 << i    # m = 2^i
    #
    # Fact: (2^(i + 1))^2 > n, so m has at least as many bits
    # as the floor of the square root of n.
    #
    # Proof: (2^(i+1))^2 = 2^(2i + 2) >= 2^(floor(log_2(n)) + 2)
    # >= 2^(ceil(log_2(n) + 1) >= 2^(log_2(n) + 1) > 2^(log_2(n)) = n. QED.
    #
    while (m << i) > n: # (m<<i) = m*(2^i) = m*m
        m >>= 1
        i -= 1
    d = n - (m << i) # d = n-m^2
    for k in xrange(i-1, -1, -1):
        j = 1 << k
        new_diff = d - (((m<<1) | j) << k) # n-(m+2^k)^2 = n-m^2-2*m*2^k-2^(2k)
        if new_diff >= 0:
            d = new_diff
            m |= j
    return m

Note that by construction, the kth bit of m << 1 is not set, so bitwise-or may be used to implement the addition of (m<<1) + (1<<k). Ultimately I have (2*m*(2**k) + 2**(2*k)) written as (((m<<1) | (1<<k)) << k), so it's three shifts and one bitwise-or (followed by a subtraction to get new_diff). Maybe there is still a more efficient way to get this? Regardless, it's far better than multiplying m*m! Compare with above:

>>> t('improved_i_sqrt(r(20000))', 'from __main__ import *', number = 5)/5.
0.10908999762373242
>>> all(improved_i_sqrt(n) == i_sqrt(n) for n in xrange(10**6))
True

Solution 5:[5]

One option would be to use the decimal module, and do it in sufficiently-precise floats:

import decimal

def isqrt(n):
    nd = decimal.Decimal(n)
    with decimal.localcontext() as ctx:
        ctx.prec = n.bit_length()
        i = int(nd.sqrt())
    if i**2 != n:
        raise ValueError('input was not a perfect square')
    return i

which I think should work:

>>> isqrt(1)
1
>>> isqrt(7**14) == 7**7
True
>>> isqrt(11**1000) == 11**500
True
>>> isqrt(11**1000+1)
Traceback (most recent call last):
  File "<ipython-input-121-e80953fb4d8e>", line 1, in <module>
    isqrt(11**1000+1)
  File "<ipython-input-100-dd91f704e2bd>", line 10, in isqrt
    raise ValueError('input was not a perfect square')
ValueError: input was not a perfect square

Solution 6:[6]

Seems like you could check like this:

if int(math.sqrt(n))**2 == n:
    print n, 'is a perfect square'

Update:

As you pointed out the above fails for large values of n. For those the following looks promising, which is an adaptation of the example C code, by Martin Guy @ UKC, June 1985, for the relatively simple looking binary numeral digit-by-digit calculation method mentioned in the Wikipedia article Methods of computing square roots:

from math import ceil, log

def isqrt(n):
    res = 0
    bit = 4**int(ceil(log(n, 4))) if n else 0  # smallest power of 4 >= the argument
    while bit:
        if n >= res + bit:
            n -= res + bit
            res = (res >> 1) + bit
        else:
            res >>= 1
        bit >>= 2
    return res

if __name__ == '__main__':
    from math import sqrt  # for comparison purposes

    for i in range(17)+[2**53, (10**100+1)**2]:
        is_perfect_sq = isqrt(i)**2 == i
        print '{:21,d}:  math.sqrt={:12,.7G}, isqrt={:10,d} {}'.format(
            i, sqrt(i), isqrt(i), '(perfect square)' if is_perfect_sq else '')

Output:

                    0:  math.sqrt=           0, isqrt=         0 (perfect square)
                    1:  math.sqrt=           1, isqrt=         1 (perfect square)
                    2:  math.sqrt=    1.414214, isqrt=         1
                    3:  math.sqrt=    1.732051, isqrt=         1
                    4:  math.sqrt=           2, isqrt=         2 (perfect square)
                    5:  math.sqrt=    2.236068, isqrt=         2
                    6:  math.sqrt=     2.44949, isqrt=         2
                    7:  math.sqrt=    2.645751, isqrt=         2
                    8:  math.sqrt=    2.828427, isqrt=         2
                    9:  math.sqrt=           3, isqrt=         3 (perfect square)
                   10:  math.sqrt=    3.162278, isqrt=         3
                   11:  math.sqrt=    3.316625, isqrt=         3
                   12:  math.sqrt=    3.464102, isqrt=         3
                   13:  math.sqrt=    3.605551, isqrt=         3
                   14:  math.sqrt=    3.741657, isqrt=         3
                   15:  math.sqrt=    3.872983, isqrt=         3
                   16:  math.sqrt=           4, isqrt=         4 (perfect square)
9,007,199,254,740,992:  math.sqrt=9.490627E+07, isqrt=94,906,265
100,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,020,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,001:  math.sqrt=      1E+100, isqrt=10,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,000,001 (perfect square)

Solution 7:[7]

Python's default math library has an integer square root function:

math.isqrt(n)

Return the integer square root of the nonnegative integer n. This is the floor of the exact square root of n, or equivalently the greatest integer a such that a² ? n.

Solution 8:[8]

The script below extracts integer square roots. It uses no divisions, only bitshifts, so it is quite fast. It uses Newton's method on the inverse square root, a technique made famous by Quake III Arena as mentioned in the Wikipedia article, Fast inverse square root.

The strategy of the algorithm to compute s = sqrt(Y) is as follows.

  1. Reduce the argument Y to y in the range [1/4, 1), i.e., y = Y/B, with 1/4 <= y < 1, where B is an even power of 2, so B = 2**(2*k) for some integer k. We want to find X, where x = X/B, and x = 1 / sqrt(y).
  2. Determine a first approximation to X using a quadratic minimax polynomial.
  3. Refine X using Newton's method.
  4. Calculate s = X*Y/(2**(3*k)).

We don't actually create fractions or perform any divisions. All the arithmetic is done with integers, and we use bit shifting to divide by various powers of B.

Range reduction lets us find a good initial approximation to feed to Newton's method. Here's a version of the 2nd degree minimax polynomial approximation to the inverse square root in the interval [1/4, 1):

Minimax poly for 1/sqrt(x)

(Sorry, I've reversed the meaning of x & y here, to conform to the usual conventions). The maximum error of this approximation is around 0.0355 ~= 1/28. Here's a graph showing the error:

Minimax poly error graph

Using this poly, our initial x starts with at least 4 or 5 bits of precision. Each round of Newton's method doubles the precision, so it doesn't take many rounds to get thousands of bits, if we want them.

""" Integer square root

    Uses no divisions, only shifts
    "Quake" style algorithm,
    i.e., Newton's method for 1 / sqrt(y)
    Uses a quadratic minimax polynomial for the first approximation

    Written by PM 2Ring 2022.01.23
"""

def int_sqrt(y):
    if y < 0:
        raise ValueError("int_sqrt arg must be >= 0, not %s" % y)
    if y < 2:
        return y

    # print("\n*", y, "*")
    # Range reduction.
    # Find k such that 1/4 <= y/b < 1, where b = 2 ** (k*2)
    j = y.bit_length()
    # Round k*2 up to the next even number
    k2 = j + (j & 1)
    # k and some useful multiples
    k = k2 >> 1
    k3 = k2 + k
    k6 = k3 << 1
    kd = k6 + 1
    # b cubed
    b3 = 1 << k6

    # Minimax approximation: x/b ~= 1 / sqrt(y/b)
    x = (((463 * y * y) >> k2) - (896 * y) + (698 << k2)) >> 8
    # print("   ", x, h)

    # Newton's method for 1 / sqrt(y/b)
    epsilon = 1 << k
    for i in range(1, 99):
        dx = x * (b3 - y * x * x) >> kd
        x += dx
        # print(f" {i}: {x} {dx}")
        if abs(dx) <= epsilon:
            break

    # s == sqrt(y)
    s = x * y >> k3
    # Adjust if too low
    ss = s + 1
    return ss if ss * ss <= y else s

def test(lo, hi, step=1):
    for y in range(lo, hi, step):
        s = int_sqrt(y)
        ss = s + 1
        s2, ss2 = s * s, ss * ss
        assert s2 <= y < ss2, (y, s2, ss2)
    print("ok")

test(0, 100000, 1)

This code is certainly slower than math.isqrt and decimal.Decimal.sqrt. Its purpose is simply to illustrate the algorithm. It would be interesting to see how fast it would be if it were implemented in C...

Here's a live version, running on the SageMathCell server. Set hi <= 0 to calculate and display the results for a single value set in lo. You can put expressions in the input boxes, eg set hi to 0 and lo to 2 * 10**100 to get sqrt(2) * 10**50.

Solution 9:[9]

Your function fails for large inputs:

In [26]: isqrt((10**100+1)**2)

ValueError: input was not a perfect square

There is a recipe on the ActiveState site which should hopefully be more reliable since it uses integer maths only. It is based on an earlier StackOverflow question: Writing your own square root function

Sources

This article follows the attribution requirements of Stack Overflow and is licensed under CC BY-SA 3.0.

Source: Stack Overflow

Solution Source
Solution 1 Nico Schlömer
Solution 2 Community
Solution 3
Solution 4 Community
Solution 5 DSM
Solution 6
Solution 7
Solution 8
Solution 9 Community

Related Questions

Selenium-Python: How do I make my WebDriverWait a few seconds before reloading a page?

Selenium-Python: How do I make my WebDriverWait a few seconds before reloading a page?

pip issue installing almost any library

pip issue installing almost any library

pip issue installing almost any library

Selenium-Python: How do I make my WebDriverWait a few seconds before reloading a page?

Selenium-Python: How do I make my WebDriverWait a few seconds before reloading a page?

A simple SMTP server (in Python)

How to specify which pytorch to import, if I have two different versions installed?

How to specify which pytorch to import, if I have two different versions installed?

A simple SMTP server (in Python)

Python: get all months in range?

Python: get all months in range?

Python: get all months in range?

How to specify which pytorch to import, if I have two different versions installed?