59 lines
1.5 KiB
Python
59 lines
1.5 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
|
||
|
|
from collections import defaultdict
|
||
|
|
from math import gcd
|
||
|
|
|
||
|
|
def sols(N):
|
||
|
|
"""yields solutions of x**2 + 3y**2 = z**2 where x + y <= 2N"""
|
||
|
|
for l in range(1, N+1, 2):
|
||
|
|
for k in range(1, N//l + 1, 2):
|
||
|
|
if gcd(l, k) != 1:
|
||
|
|
continue
|
||
|
|
x = abs(3*l**2 - k**2) // 2
|
||
|
|
y = l * k
|
||
|
|
z = (3*l**2 + k**2) // 2
|
||
|
|
for n in range(1, (2*N) // (x+y) + 1):
|
||
|
|
yield (n*x, n*y, n*z)
|
||
|
|
|
||
|
|
for l in range(1, N//2 + 1):
|
||
|
|
for k in range(1 + (l%2), N//(2*l) + 1, 2):
|
||
|
|
if gcd(l, k) != 1:
|
||
|
|
continue
|
||
|
|
x = 2 * abs(3*l**2 - k**2)
|
||
|
|
y = 4 * l * k
|
||
|
|
z = 2 * (3*l**2 + k**2)
|
||
|
|
if x + y > 2*N:
|
||
|
|
continue
|
||
|
|
for n in range(1, (2*N) // (x+y) + 1):
|
||
|
|
yield (n*x, n*y, n*z)
|
||
|
|
|
||
|
|
def pq_gen(N):
|
||
|
|
"""yields solutions of c**2 = p**2 + q**2 + pq where c is int and p + q <= N"""
|
||
|
|
for x, y, z in sols(N):
|
||
|
|
p = y
|
||
|
|
if x <= p or (x - p) % 2 != 0:
|
||
|
|
continue
|
||
|
|
q = (x - p) // 2
|
||
|
|
yield p, q
|
||
|
|
|
||
|
|
|
||
|
|
def find_pqr(N):
|
||
|
|
pqs = defaultdict(lambda : set())
|
||
|
|
for p, q in pq_gen(N-1):
|
||
|
|
pqs[p].add(q)
|
||
|
|
pqs[q].add(p)
|
||
|
|
sums = set()
|
||
|
|
for p, qrs in dict(pqs).items():
|
||
|
|
for q in qrs:
|
||
|
|
for r in qrs:
|
||
|
|
if p + q + r > N or r not in pqs[q]:
|
||
|
|
continue
|
||
|
|
sums.add(p + q + r)
|
||
|
|
return sum(sums)
|
||
|
|
|
||
|
|
|
||
|
|
N = 120 * 1000
|
||
|
|
|
||
|
|
print(find_pqr(N))
|
||
|
|
|