# Python optimization search pythagorean triplets

I have an optimization problem. Be it about code, algorithm or maths, that, i don't know. My goal is to find pythagorean triplets satisfying a = b + 1 or a = b - 1. m and n are the triplets generators. The following code produces the desired output, but it's not fast enough from about the 11th or 12th item.

#! /usr/bin/env python3

from time import perf_counter
import numpy as np
from numba import njit, prange
import json, os

@njit
def isPythagoreanTriplet(x:int, y:int) -> bool:
"""
Test if a and b are a valid basis for a pythagorean triplet
"""
return (abs(2*x - y) == 1)

@njit
def findNextPythagoreanTriplet(n:int)->tuple[int]:
"""
Self-explanatory function name
"""
m = 4 * n
while True:
for n in prange(1, m):
m_sq = m * m
n_sq = n * n
a = m_sq - n_sq
b = 2 * m * n
c = m_sq + n_sq
if isPythagoreanTriplet(a, b):
return(2*a, b, c, m, n)
elif isPythagoreanTriplet(b, a):
return(a, 2*b, c, m, n)
m+= 1

def pythagoreanTriplets(compteur:int, n:int=1)->None :
"""
Find compteur pythagorean triplets
"""
while compteur :
a, b, c, m, n = findNextPythagoreanTriplet(n)
store(a, b, c, m, n)
n = m
compteur-= 1

def store(a,b,c, m, n):
"""
Store a pythagorean triplet and its generators in a text file
"""
with open("results.txt", "a") as file_out:
file_out.write(f"{a} {b} {c} {m} {n}\n")

def init_n()->tuple[int]:
"""
Check if result text file exists to start from last stop for n
"""
nb_lignes = 0
n = 1
if os.path.exists("results.txt"):
with open("results.txt", "r") as file_in:
for line in lines:
line = line.strip()
if line:
nb_lignes+= 1
_, _, _, m, n = line.split()
n = m

return (nb_lignes, int(n))

# Driver Code
if __name__ == '__main__' :
nb_items = 12
start = perf_counter()
nb_lignes, n = init_n()
nb_items-= nb_lignes
print(f"Nb items restants : {nb_items}")
pythagoreanTriplets(nb_items, n)
stop = perf_counter()
duree = stop - start
print(f"Durée : {duree:.2f}s")

• Micro-review: PEP-8 recommends snake_case for function names. Commented Dec 8, 2023 at 14:43

I think you're coming at this with the wrong algorithm.

Given a right triangle with sides a, b and (hypotenuse) a+1, we have:

$$a² + b² = (a+1)²$$

So

$$b² = (a+1)² - a²$$ $$= a² + 2a + 1 - a²$$ $$= 2a + 1$$

From this, we can deduce that for any odd length for b (b > 1), we can find a corresponding a and therefore the other two sides.

Using that algorithm, I have this generator (lacking docstring and type annotations):

def generate_triplets():
b = 1
while True:
b += 2
a = b * b // 2
yield (a, b, a+1)


Simple demo:

import itertools
for triangle in itertools.islice(generate_triplets(), 50):
print(triangle)