Python optimization search pythagorean triplets

Question

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:
 lines = file_in.readlines()
 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") 

• Toby Speight: Micro-review: PEP-8 recommends snake_case for function names.

Answer 1

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:

$$ a2 + b2 = (a+1)2 $$

So

$$ b2 = (a+1)2 - a2 $$ $$ = a2 + 2a + 1 - a2 $$ $$ = 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)