Python中的整数平方根

Question

在Python或者标准库里，有没有整数平方根的功能？我希望它能返回一个整数，并且如果输入的数字不是完全平方数，就能抛出一个异常。

我试过用这段代码：

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

但是我觉得这段代码不太好看，而且对于大整数我也不太放心。我可以通过遍历平方数来找到结果，如果超过了目标值就放弃，但我觉得这样做可能会比较慢。而且，这种功能应该已经有人实现过了吧？

---

_{另见： 检查一个数字是否是完全平方数.}

Answer 1

抱歉回复得很晚；我刚好看到这个页面。如果将来有人访问这个页面，python模块gmpy2是为了处理非常大的输入而设计的，其中包括一个整数平方根的功能。

举个例子：

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

当然，所有的结果都会带有“mpz”这个标签，但mpz和int是可以兼容的：

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

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

可以查看我其他的回答，里面讨论了这个方法的性能与其他一些回答的比较。

下载地址：https://code.google.com/p/gmpy/

Answer 2

更新： Python 3.8 版本的标准库中有一个 math.isqrt 函数！

我对这里每个（正确的）函数进行了性能测试，测试了小范围（0…2²²）和大范围（2⁵⁰⁰⁰¹）的输入。结果显示，在这两种情况下，表现最好的分别是 gmpy2.isqrt（由 mathmandan 提出的） 排名第一，其次是 Python 3.8 的 math.isqrt，排名第三的是 NPE 提供的 ActiveState 方案。ActiveState 的方案有很多除法运算，可以用位移操作替代，这样会稍微快一点（但仍然不如原生函数快）：

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")

性能测试结果：

• gmpy2.isqrt()（mathmandan）: 小范围 0.08 微秒，大范围 0.07 毫秒
• int(gmpy2.isqrt())*: 小范围 0.3 微秒，大范围 0.07 毫秒
• Python 3.8 math.isqrt: 小范围 0.13 微秒，大范围 0.9 毫秒
• ActiveState（如上优化）: 小范围 0.6 微秒，大范围 17.0 毫秒
• ActiveState（NPE）: 小范围 1.0 微秒，大范围 17.3 毫秒
• castlebravo 手动实现: 小范围 4 微秒，大范围 80 毫秒
• mathmandan 改进版: 小范围 2.7 微秒，大范围 120 毫秒
• martineau（带有 这个修正）: 小范围 2.3 微秒，大范围 140 毫秒
• nibot: 小范围 8 微秒，大范围 1000 毫秒
• mathmandan: 小范围 1.8 微秒，大范围 2200 毫秒
• castlebravo 牛顿法: 小范围 1.5 微秒，大范围 19000 毫秒
• user448810: 小范围 1.4 微秒，大范围 20000 毫秒

(* 因为 gmpy2.isqrt 返回的是一个 gmpy2.mpz 对象，它的行为大部分和 int 相似，但并不完全相同，所以在某些情况下你可能需要把它转换回 int。)

Answer 3

注意： 从Python 3.8开始，标准库中新增了一个叫math.isqrt的功能。

牛顿法在处理整数时效果很好：

def isqrt(n):
    x = n
    y = (x + 1) // 2
    while y < x:
        x = y
        y = (x + n // x) // 2
    return x

这个方法会返回一个最大的整数x，使得x * x 不超过n。如果你想检查这个结果是否正好是平方根，可以直接进行乘法运算，看看n是否是一个完全平方数。

我在我的博客上讨论了这个算法，以及另外三个计算平方根的算法。