540. 本原毕达哥拉斯三元组计数(Counting primitive Pythagorean triples)

毕达哥拉斯三元组包含有三个整数\(a,b,c\)且满足等式\(a^2+b^2=c^2\)，当\(a,b,c\)互素时，这个三元组被称为本原的。记\(P(n)\)为满足\(a<b<c \le n\)的本原毕达哥拉斯三元组的数目。例如\(P(20)=3\)，这三个三元组是：\((3,4,5),(5,12,13)\)和\((8,15,17)\)。已知\(P(10^6)=159139\)，求\(P(3141592653589793)\)。

分析：针对这道题我虽然尝试了很多方法，最终也给出了正确答案，但是程序的运行时间达到了36分钟，显然不符合欧拉计划网站的标准，所以不算完全解决了这个问题。但是也不妨记录一下思路，等到将来有更好的方法再做改进。题目要求满足斜边长度不超过\(n\)的本原毕达哥拉斯三元组的个数，首先我们需要知道这些三元组如何产生出来。我们可以使用在以前的题目中就介绍过的欧几里德公式，即当\(gcd(x,y)=1,x>y\ge1\)且\(x,y\)不同为奇数时，本原毕达哥拉斯三元组可以使用以下公式产生：

因此我们的问题可以转化成求满足以下条件的\((x,y)\)对的个数：

• 斜边长度不能超过\(n\)，即\(x^2+y^2\le n, where\ x>y\ge1\)；
• \(x,y\)是互质关系，即\(gcd(x,y)=1\)
• \(x,y\)不同时为奇数，即当\(x\)是偶数时，\(y\)可奇可偶，而当\(x\)是奇数时，\(y\)只能是偶数。

我们可以用图示让上述要求看得更清楚，针对第一个条件即\(x^2+y^2\le n=(\sqrt{n})^2\)，即\(x,y\)要在圆心在原点，半径为\(\sqrt{n}\)的圆内及圆上，同时要求\(x>y\)，则这些点还要在\(y=x\)的直线下方。在下面的图示中，我们要求三元组的斜边长度不超过100，即\(x,y\)都在半径为10原圆内或圆上及位于绿色直线下方，即位于圆\(x^2+y^2=10\)，直线\(y=x\)以及正横轴三者围住的区域，而\((x,y)\)正对应这个区域内的格点(即横纵坐标均为整数的点)。

需要注意的是，橙色直线\(BC\)把这个区域分成了左右两个部分，在左边绿色直线位于圆周下方，则起作用的约束条件是\(x>y\)；而在右边绿色直线位于圆周上方，这时起作用的约束条件是\(x^2+y^2\le n\)，即\(y\le \sqrt{n-x^2}\)。

在以上区域中，我们找到了16个格点，也即斜边长度不超过100的本原毕达哥拉斯三元组共有16个。比如当\(x=5\)时，我们首先看\((5,1)\)这个格点，因为条件\(x,y\)不能同时为奇数，所以这个格点排除，同理可以把\((5,3)\)排除；再看\((5,2)\)这个格点，满足\(gcd(5,2)=1\)，则一个符合要求的格点，同理\((5,4)\)也符合。再往上就不满足\(x>y\)的要求了，所以格点\((5,5)\)以及的格点都排除点。基于以上思路，我们可以把满足要求的格点分为四种情况：

• 当\(x<\sqrt{n/2}\)时，此时\(y<x\)：\(x\)为偶数，则要求小于\(x\)的数中与\(x\)互质的数的个数，即求欧拉函数\(\varphi(x)\)；\(x\)为奇数，也要求小于\(x\)的数中与\(x\)互质的数的个数，但\(y\)不能也为奇数，所以是求\(\varphi(x)/2\)；

• 当\(x\ge\sqrt{n/2}\)时，此时\(y\le \sqrt{n-x^2}\)：\(x\)为偶数，则要求小于等于\(\sqrt{n-x^2}\)的数中与\(x\)互质的数的个数，即求广义欧拉函数\(\phi \left( x,\left\lfloor \sqrt{n-x^{2}}\right\rfloor \right)\)；\(x\)为奇数，则要求小于等于\(\sqrt{n-x^2}\)的数中与\(x\)互质的数的个数但\(y\)不能也为奇数，即求广义欧拉函数\(\phi \left( x,\left\lfloor \frac{\sqrt{n-x^{2}}}{2}\right\rfloor \right)\)；

欧拉函数\(\varphi(x)\)可以比较容易求得，但考虑到计算中需要大量的欧拉函数值，所以可以用筛法预先计算出这次值，这样比每次需要时再做计算效率要更高一些，具体实现可以参见代码中的\(euler\_phi\_sieve()\)函数。而对于广义欧拉函数，使用莫比乌斯函数，我们也可以计算出来：

为加快计算，我们可以使用\(spymp\)模块中的\(mobiusrange()\)函数预先计算出所需的莫比乌斯函数值。完整代码如下：

# time cost = 35min 37s

from sympy import sieve,divisors

def euler_phi_sieve(n):
    arr = list(range(n+1))
    for i in range(2,n+1):
        if arr[i] == i:
            for j in range(i,n+1,i):
                arr[j] -= arr[j]//i
    return arr

def phi(m,n,premob):
    res = 0
    for i in [x for x in divisors(m) if x<=n]:
        res += (premob[i-1]*(n//i))
    return res

def main(N=3141592653589793):
    res = 0
    u = int((N/2)**0.5)+1
    prephi = euler_phi_sieve(u)
    premob = list(sieve.mobiusrange(1,u))
    for m in range(2,u):
        if m % 2 == 0:
            res += prephi[m]
        else:
            res += (prephi[m]//2)
    for m in range(u,int(N**0.5)+1):
        v = int((N-m**2)**0.5)
        if m % 2 == 0:
            res += phi(m,v,premob)
        else:
            res += phi(m,v//2,premob)
    return res