Beginner Exercism • python

Pythagorean Triplet

Challenge Overview

Given an integer N, find all Pythagorean triplets for which a + b + c = N.

Introduction

You are an accomplished problem-solver, known for your ability to tackle the most challenging mathematical puzzles. One evening, you receive an urgent letter from an inventor called the Triangle Tinkerer, who is working on a groundbreaking new project. The letter reads:

Dear Mathematician,

I need your help. I am designing a device that relies on the unique properties of Pythagorean triplets — sets of three integers that satisfy the equation a² + b² = c². This device will revolutionize navigation, but for it to work, I must program it with every possible triplet where the sum of a, b, and c equals a specific number, N. Calculating these triplets by hand would take me years, but I hear you are more than up to the task.

Time is of the essence. The future of my invention — and perhaps even the future of mathematical innovation — rests on your ability to solve this problem.

Motivated by the importance of the task, you set out to find all Pythagorean triplets that satisfy the condition. Your work could have far-reaching implications, unlocking new possibilities in science and engineering. Can you rise to the challenge and make history?

Instructions

Description

A Pythagorean triplet is a set of three natural numbers, {a, b, c}, for which,

a² + b² = c²

and such that,

a < b < c

For example,

3² + 4² = 5².

Given an input integer N, find all Pythagorean triplets for which a + b + c = N.

For example, with N = 1000, there is exactly one Pythagorean triplet for which a + b + c = 1000: {200, 375, 425}.

Dig Deeper

Cubic

Cubic-time triply-nested loops

def triplets_with_sum(n):
    triplets = []
    for a in range(1, n + 1):
        for b in range(a + 1, n + 1):
            for c in range(b + 1, n + 1):
                if a**2 + b**2 == c**2 and a + b + c == n:
                    triplets.append([a, b, c])
    return triplets

The strategy in this case is to scan through all three variables and test them in the innermost loop. But the innermost loop will test all variables every time the enclosing loop iterates. And the enclosing loop up will iterate through all of its range every time the outermost loop iterates.

So the ‘work’ this code has to do for each additional value in range(1, n+1) is n**3.

This gives code that is simple, clear, …and so slow as to be useless for all but the smallest values of n.

We could tighten up the bounds on loop variables: for example, a is the smallest integer of a triplet that sums to n, so inevitably a < n //3.

However, this is not nearly enough to rescue an inappropriate algorithm.

Quadratic

Quadratic-time doubly-nested loops

def triplets_with_sum(n):
    triplets = []
    for a in range(1, n + 1):
        for b in range(a + 1, n + 1):
            c = n - a - b
            if a ** 2 + b ** 2 == c ** 2:
                triplets.append([a, b, c])
    return triplets

Because a + b + c == n, we only loop over a and b. The third variable c is then predictable.

The above code loops over the full range of both variables. This means the ‘work’ this code has to do for each additional value in range(1, n+1) is n**2. We know enough about the problems to tighten this up a bit.

For example:

The smallest pythagorean is (famously) [3, 4, 5], so a >= 3
a + b == n - c and a <= b, so a <= (n - c) // 2

We can also avoid, to some extent, repeating the same multiplication many times. This gets us to the code below:

def triplets_with_sum(n):
    result = []
    for c in range(5, n - 1):
        c_sq = c * c
        for a in range(3, (n - c + 1) // 2):
            b = n - a - c
            if a < b < c and a * a + b * b == c_sq:
                result.append([a, b, c])
    return result

We could have done a bit better. Though not stated in the problem description, a + b > c, otherwise they could not form a triangle.

This means that c <= n // 2, reducing the iterations in the outer loop.

However, these optimizations only reduce the run time by a small factor. They do almost nothing to make the algorithm scale to large n.

Linear Loop

Linear-time algorithms with no nested loops

The key point with this approach is that we only loop over the variable c in a specific range. Some mathematical analysis (essentially, solving simultaneous equations) then allows us to find valid values of a and b.

Other than that, the code syntax below is fairly mainstream across programming languages. A related approach instead loops over a.

from math import sqrt

def triplets_with_sum(n):
    N = float(n)
    triplets = []
    for c in range(int(N / 2) - 1, int((sqrt(2) - 1) * N), -1):
        D = sqrt(c ** 2 - N ** 2 + 2 * N * c)
        if D == int(D):
            triplets.append([int((N - c - D) / 2), int((N - c + D) / 2), c])
    return triplets

This second code example has no explicit for loop (in Python syntax), but the generator-expression and the list-comprehension both translate to FOR_ITER at the bytecode level. So this solution is essentially the same as the first, written in a more ‘Pythonic’ syntax — but that syntax does incur a small overhead in performance. The performance hit is likely due to the extra instructions in the bytecode used to manage the generator-expression (pausing the loop, resuming the loop, yielding results) and then calling or unpacking the generator in the list comprehension. However, you would have to carefully profile the code to really determine the slowdown.

With all that said, using a generator or generator-expression with or without a list-comprehension might be a better strategy if your code needs to process a very large number of triplets, as it avoids storing all the results in memory until they need to be returned. Using a generator or generator-expression by itself can also nicely set up a scenario where results are “streamed” or emitted ‘on demand’ for another part of the program or application.

For more details on what these two solutions look like at the byte code level, take a look at Pythons dis module.

def triplets_with_sum(n):
    def calculate_medium(small):
        return (n ** 2 - 2 * n * small) / (2 * (n - small))

    two_sides = ((int(medium), small) for small in range(3, n // 3)
                 if small < (medium := calculate_medium(small))
                 and medium.is_integer())

    return [[small, medium, (medium ** 2 + small ** 2) ** 0.5]
            for medium, small in two_sides]

Some implementation details to notice:

Nested functions, with the inner function able to reference variables such as n passed into the outer function.
The generator expression creates two_sides as a lazily-evaluated iterator (smaller memory footprint)
The walrus operator := is new in Python 3.8.
The is_integer() method replaces if D == int(D).
Using ** 0.5 to calculate the square roots avoids a math import.

Source: Exercism python/pythagorean-triplet