9
\$\begingroup\$

How can I make this code run faster:

for a in range(1,1001):
    for b in range(1, 1001):
        for c in range(1, 1001):
            if pow(a, 2) + pow(b, 2) == pow(c, 2):
                print(str(a) + "," + str(b) + "," + str(c))
\$\endgroup\$
  • 5
    \$\begingroup\$ You can generate a tree of all pythagorean triples using some simple linear algebra. \$\endgroup\$ – Phylogenesis Oct 20 at 0:25
  • 1
    \$\begingroup\$ @Phylogenesis Very nice, used that in another answer. \$\endgroup\$ – Stefan Pochmann Oct 20 at 16:35
  • 1
    \$\begingroup\$ Try replacing pow(X,2 with (X * X). Usually, a single multiplication is faster than a function call. \$\endgroup\$ – Thomas Matthews Oct 20 at 16:44
  • \$\begingroup\$ The naive approach is O(N^3), since each for-loop runs 1..1001. Also, you can contrain the start and termination values on your c-loop by applying a,b < c (WLOG let a,b be the legs and c the hypotenuse). But even if you don't code the tree-based approach, you can reduce the order enormously by generating all distinct positive-integer candidates (m,n) in Euclid's solution a = k(m^2 - n^2), b = k(2mn), c = k(m^2 + n^2). Now we only need to test up to c = 1000 > m^2 i.e. m < sqrt(1000) = 31 @StefanPochmann's answer \$\endgroup\$ – smci Oct 20 at 22:18
  • \$\begingroup\$ By the way, using f-strings in printing is clearer: print(f'{a},{b},{c}') \$\endgroup\$ – smci Oct 20 at 23:04
20
\$\begingroup\$

Some optimizations and style suggestions:

  • After finding a solution you can break:
    for a in range(1,1001):
        for b in range(1, 1001):
            for c in range(1, 1001):
                if pow(a, 2) + pow(b, 2) == pow(c, 2):
                    print(str(a) + "," + str(b) + "," + str(c))
                    break
    
  • Use ** which is faster than pow, or just multiply for itself a*a.
  • Use Python formatter to print the result: print(f"{a},{b},{c}").
  • Calculate c as \$c=sqrt(a^2+b^2)\$:
    for a in range(1,1001):
        for b in range(1, 1001):
            c = int(math.sqrt(a ** 2 + b ** 2))
            if a ** 2 + b ** 2 == c ** 2 and c < 1001:
                print(f"{a},{b},{c}")
    
    The solution now takes \$O(n^2)\$ instead of \$O(n^3)\$.
  • Instead of checking if a ** 2 + b ** 2 == c ** 2:, it's enough verifying that c is an integer:
    for a in range(1,1001):
      for b in range(1, 1001):
          c = math.sqrt(a ** 2 + b ** 2)
          if c.is_integer() and c < 1001:
              print(f"{a},{b},{int(c)}")
    
  • As already said, you can also start the second for-loop from a to avoid duplicated solutions.
  • Put everything into a function:
    def triplets(n):
      for a in range(1, n):
          for b in range(a, n):
              c = math.sqrt(a * a + b * b)
              if c.is_integer() and c <= n:
                  print(f"{a},{b},{int(c)}")
    triplets(1000)
    

Runtime on my machine:

Original: 868.27 seconds (~15 minutes)
Improved: 0.27 seconds

EDIT:

Since this question got a lot of attention I wanted to add a couple of notes:

  1. This early answer was for OP's original problem that I interpreted as "find all the triplets in a reasonable amount of time".
  2. There are definitely more efficient (and advanced) solutions than this one. If you are interested in finding out more, have a look at other excellent answers on this thread.
  3. As noted in the comments, the runtime in my answer is a rough calculation. Find a better benchmark on @Stefan's answer.
| improve this answer | |
\$\endgroup\$
  • 2
    \$\begingroup\$ Thank you so much! \$\endgroup\$ – Chezhiiyan Sabapathy Oct 19 at 9:08
  • 1
    \$\begingroup\$ I was very much in a rush when I wrote my answer. This is orders of magnitude better. \$\endgroup\$ – Jörg W Mittag Oct 19 at 15:57
  • 2
    \$\begingroup\$ For the runtimes, how did you measure? Asking mainly because you print the results. \$\endgroup\$ – superb rain Oct 19 at 16:31
  • 2
    \$\begingroup\$ You can still improve (slightly) the inner loop by running from a up to sqrt(n**2-a**2) and discard the c<=n test. \$\endgroup\$ – Thomas Baruchel Oct 19 at 16:44
  • 1
    \$\begingroup\$ @superbrain I used time.time() around triplets(1000), that prints to the console. Then I wrapped OP's solution into a function and calculated the runtime the same way. It was just to have a rough idea, a better benchmark is in @Stefan's answer. \$\endgroup\$ – Marc Oct 20 at 2:04
16
\$\begingroup\$

My "review" will have to be "If you really want it fast, you need a completely different approach". The following ~ O(N log N) approach is about 680 times faster than Marc's accepted solution for N=1000:

from math import isqrt, gcd

def triplets(N):
    for m in range(isqrt(N-1)+1):
        for n in range(1+m%2, min(m, isqrt(N-m*m)+1), 2):
            if gcd(m, n) > 1:
                continue
            a = m*m - n*n
            b = 2*m*n
            c = m*m + n*n
            for k in range(1, N//c+1):
                yield k*a, k*b, k*c

This uses Euclid's formula.

Benchmark results for N=1000:

Stefan    Marc
0.24 ms   165.51 ms
0.24 ms   165.25 ms
0.24 ms   161.33 ms

Benchmark results for N=2000, where it's already about 1200 times faster than the accepted solution:

Stefan    Marc      
0.52 ms   654.72 ms   
0.58 ms   689.10 ms   
0.53 ms   662.19 ms   

Benchmark code:

from math import isqrt, gcd
import math
from timeit import repeat
from collections import deque

def triplets_Stefan(N):
    for m in range(isqrt(N-1)+1):
        for n in range(1+m%2, min(m, isqrt(N-m*m)+1), 2):
            if gcd(m, n) > 1:
                continue
            a = m*m - n*n
            b = 2*m*n
            c = m*m + n*n
            for k in range(1, N//c+1):
                yield k*a, k*b, k*c

def triplets_Marc(n):
  for a in range(1, n):
      for b in range(a, n):
          c = math.sqrt(a * a + b * b)
          if c.is_integer() and c <= n:
              yield a, b, int(c)

n = 2000
expect = sorted(map(sorted, triplets_Marc(n)))
result = sorted(map(sorted, triplets_Stefan(n)))
print(expect == result)

funcs = [
    (10**3, triplets_Stefan),
    (10**0, triplets_Marc),
    ]

for _, func in funcs:
    print(func.__name__.removeprefix('triplets_').ljust(10), end='')
print()

for _ in range(3):
    for number, func in funcs:
        t = min(repeat(lambda: deque(func(n), 0), number=number)) / number
        print('%.2f ms   ' % (t * 1e3), end='')
    print()

About runtime complexity: Looks like around O(N log N). See the comments. And if I try larger and larger N = 2**e and divide the times by N log N, they remain fairly constant:

>>> from timeit import repeat
>>> from collections import deque
>>> for e in range(10, 25):
        N = 2**e
        t = min(repeat(lambda: deque(triplets(N), 0), number=1))
        print(e, t / (N * e))

10 5.312499999909903e-08
11 3.3176491483337275e-08
12 2.3059082032705902e-08
13 3.789156400398811e-08
14 1.95251464847414e-08
15 1.9453328450880215e-08
16 1.9563865661601648e-08
17 1.9452756993864518e-08
18 1.973256005180039e-08
19 2.0924497905514347e-08
20 2.1869220733644352e-08
21 2.1237255278089392e-08
22 2.0788834311744357e-08
23 2.1097218990325713e-08
24 2.1043718606233202e-08

Also see the comments.

| improve this answer | |
\$\endgroup\$
  • 14
    \$\begingroup\$ Yay 2200 year old maths! \$\endgroup\$ – Jörg W Mittag Oct 19 at 17:38
  • 2
    \$\begingroup\$ @smci I don't know. But yes, looks like around O(N log N). \$\endgroup\$ – Stefan Pochmann Oct 20 at 22:35
  • 1
    \$\begingroup\$ @smci I added some timings for larger N at the end (up to \$N=2^{24}\$). Don't know about the other gcds, and I'm not sure "pruning" is the right word. To me that sounds like a performance optimization and that's how you're treating it, but its real purpose here is to avoid duplicates in the output. \$\endgroup\$ – Stefan Pochmann Oct 20 at 23:11
  • 1
    \$\begingroup\$ @smci Meh, I don't think we need a graph to know that ~O(N log N) is much faster than O(N^2) and O(N^3) :-) \$\endgroup\$ – Stefan Pochmann Oct 20 at 23:36
  • 1
    \$\begingroup\$ It would be really interesting to see how this compares on different Python implementations. The OP's extremely inefficient version runs in ~15 minutes on my laptop using CPython 3.7.3 but in only 22 seconds, i.e. 40x(!!!) faster, on GraalPython 20.2.0. And this is actually still unfair towards GraalPython because I did not account for JVM startup time, JIT warmup, etc. Even more interesting: GraalPython used 200% CPU, so somehow, in some way, shape, or form, it was using multiple threads. No idea how and why and what, though. It can't be the GC since there's not garbage. It can't be the JIT … \$\endgroup\$ – Jörg W Mittag Oct 21 at 11:24
8
\$\begingroup\$

There are some obvious optimizations you can do:

  • Eliminate duplicate checks: you are checking each possible combination twice. For example, you are checking both a = 2, b = 3 and a = 3, b = 2. The very first two lines your program prints are 3,4,5 and 4,3,5!
  • Eliminate impossible checks: you are checking, for example, 1000² + 1000² == 1².
  • Eliminate duplicate computations: you are computing the square of the same numbers over and over again.
  • Print less: printing to the console is sloooooooooooow. Collect the results in a data structure and only print them once.

Something like this:

def triplets():
    squares = [pow(n, 2) for n in range(0, 1001)]

    for a in range(1, 1001):
        for b in range(a, 1001):
            for c in range(b, 1001):
                if squares[a] + squares[b] == squares[c]:
                    yield a, b, c


print(list(triplets()))
```
| improve this answer | |
\$\endgroup\$
  • 1
    \$\begingroup\$ Given that you've precomputed all the relevant squares, you can just do list membership to see whether the result is a square, making the code about 4 times faster: Try it online! \$\endgroup\$ – Neil Oct 21 at 9:59
5
\$\begingroup\$

First, I don't know Python, so please don't look to me as setting a stylistic or idiomatic example here. But I think that there are some things that are universal. In particular, try to move calculations out of loops. So in your original (although the same advice applies to all the answers posted so far in some way):

for a in range(1, 1001):
    square_a = a * a
    for b in range(1, 1001):
        square_c = square_a + b * b
        for c in range(1, 1001):
            if square_c == c * c:

It is possible that the Python compiler or interpreter will do this for you, pulling the invariant calculations out of the loops. But if you do it explicitly, then you know that it will be done.

You can use the benchmarking techniques in Stefan Pochmann's answer to test if it helps.

| improve this answer | |
\$\endgroup\$
  • 1
    \$\begingroup\$ The most widely used Python implementation, CPython, is a pure interpreter. It won't hoist a*a for you. And might not optimize pow(x, 2) into a simply multiply. BTW, your way makes the optimization of simply checking if square_c is a perfect square between 1 and 1001 more obvious. Like tmp = sqrt(square_c); if int(tmp) == tmp: instead of looping. (or something like that; IDK Python either.) \$\endgroup\$ – Peter Cordes Oct 21 at 3:37
5
\$\begingroup\$

Trees of primitive Pythagorean triples are great. Here's a solution using such a tree:

def triplets(N):
    mns = [(2, 1)]
    for m, n in mns:
        c = m*m + n*n
        if c <= N:
            a = m*m - n*n
            b = 2 * m * n
            for k in range(1, N//c+1):
                yield k*a, k*b, k*c
            mns += (2*m-n, m), (2*m+n, m), (m+2*n, n)

And here's one using a heap to produce triples in increasing order of c:

from heapq import heappush, heappop

def triplets(N=float('inf')):
    heap = []
    def push(m, n, k=1):
        kc = k * (m*m + n*n)
        if kc <= N:
            heappush(heap, (kc, m, n, k))
    push(2, 1)
    while heap:
        kc, m, n, k = heappop(heap)
        a = m*m - n*n
        b = 2 * m * n
        yield k*a, k*b, kc
        push(m, n, k+1)
        if k == 1:
            push(2*m-n, m)
            push(2*m+n, m)
            push(m+2*n, n)

A node in the primitive triple tree just needs its m and n (from which a, b and c are computed). I instead store tuples (kc, m, n, k) in a heap, where k is the multiplier for the triple and c is the primitive triple's c so that kc is the multiplied triple's c. This way I get all triples in order of increasing (k-multiplied) c. The tree structure makes the expansion of a triple to larger triples really easy and natural. I had tried to do something like this with my loops-solution but had trouble. Also note that I don't need any ugly sqrt-limit calculations, don't need a gcd-check, and don't need to explicitly make sure m+n is odd (all of which I have in my other answer's solution).

Demo:

>>> for a, b, c in triplets():
        print(a, b, c)
        
3 4 5
6 8 10
5 12 13
9 12 15
15 8 17
12 16 20
...
(I stopped it here)

So if you want the triples up to a certain limit N, you can provide it as argument, or you can just read from the infinite iterator and stop when you exceed the limit or when you've had enough or whatever. For example, the millionth triple has c=531852:

>>> from itertools import islice
>>> next(islice(triplets(), 10**6-1, None))
(116748, 518880, 531852)

This took about three seconds.

Benchmarks with my other answer's "loops" solution, the unordered "tree1" solution and the ordered-by-c "tree2" solution:

N = 1,000
loops     tree1     tree2     
0.25 ms   0.30 ms   1.14 ms   
0.25 ms   0.31 ms   1.18 ms   
0.25 ms   0.32 ms   1.15 ms   

N = 2,000
loops     tree1     tree2     
0.53 ms   0.61 ms   2.64 ms   
0.52 ms   0.60 ms   2.66 ms   
0.51 ms   0.60 ms   2.54 ms   

N = 1,000,000
loops     tree1     tree2     
0.46 s    0.52 s    6.02 s   
0.47 s    0.53 s    6.04 s   
0.45 s    0.53 s    6.08 s   

Thanks to @Phylogenesis for pointing these trees out.

| improve this answer | |
\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.