Python 3.4. Code that searches a sequence of integers for those that are the result of the sum of two squares

Question:

As a Python practice, I write a code that searches a sequence for which integers are the result of the sum of two (or more) squares. It works well but I have the following doubts.

  1. I start with a generator and an empty list. The numbers are compared with the content of the list and added to it if they meet certain conditions. Am I correct if this is more efficient than initially creating a list (instead of the generator)? Or does it depend on whether in the first case I prefer more available memory or, in the second, less data processing? (I do not know if I explained well).

  2. How could the code be improved? I'd like to see if someone more advanced (I'm a hobbyist and self-taught) can come up with a noticeably simpler solution. I must have done something stupid.

The program is as follows:

def gen_raices(min, max):
    lista = []

    # Comprueba entre tres opciones:
    # op1 = El entero no está en lista.
    # op2 = Está en lista pero de suma distintas.
    # op3 = Está en lista pero con misma suma.
    def comprueba (r1, r2):
        for d in range(len(lista)):
            if lista[d][0] == r1 + r2:
                for t in range(1,len(lista[d])):
                    if r1 in lista[d][t]:
                        return 3
                return (2, d)
        return 1

    # Añade a la lista las sumas de raices.
    e1 = 1
    while e1 ** 2 <= max:
        for e2 in range(1, max):
            r1 = e1 ** 2
            r2 = e2 ** 2
            if r1 + r2 <= max and r1 + r2 >= min:
                comp = comprueba(r1, r2)
                if comp == 1:
                    lista.append([r1 + r2, (r1, r2)])
                elif comp == 3:
                    continue
                else:
                    lista[comp[1]].append((r1, r2))
        e1 += 1
    lista.sort()
    return lista

print(gen_raices(1,100))  

the result is:

[[2, (1, 1)], [5, (1, 4)], [8, (4, 4)], [10, (1, 9)], [13, (4, 9)], [17, (1, 16)], [18, (9, 9)], [20, (4, 16)], [25, (9, 16)], [26, (1, 25)], [29, (4, 25)], [32, (16, 16)], [34, (9, 25)], [37, (1, 36)], [40, (4, 36)], [41, (16, 25)], [45, (9, 36)], [50, (1, 49), (25, 25)], [52, (16, 36)], [53, (4, 49)], [58, (9, 49)], [61, (25, 36)], [65, (1, 64), (16, 49)], [68, (4, 64)], [72, (36, 36)], [73, (9, 64)], [74, (25, 49)], [80, (16, 64)], [82, (1, 81)], [85, (4, 81), (36, 49)], [89, (25, 64)], [90, (9, 81)], [97, (16, 81)], [98, (49, 49)], [100, (36, 64)]]

Answer:

Preliminary note

First point out that in your question you mention the term "generator", but your code doesn't really use generators. This word has a very specific meaning in python. A generator is a function that contains a yield statement, which somehow return as a return , but without terminating the function, which can be resumed later and continued on the line following the yield .

A generator could be used in your case so that, every time you call it, it "generates" the next sum of squares (instead of saving them all in a list and at the end return the list), but that is not what you are. making.

Improvements in your code

The code is correct and as you say it works, but there are many possible improvements

Most appropriate data types

The types you use to solve the problem are basically lists. Python has others more appropriate to speed up searches.

For example, your list stores for each number, its decomposition as a sum of squares. Thus, an element of the list has, for example, the values [65, (1, 64), (16, 49)] , which indicates that 65 is decomposed into the sum 1+64 and also 16+49 .

A more efficient way would be that instead of a list it was a dictionary , whose key is each number, and each value the list of possible decompositions into squares. So your dictionary would have (among others) the entry { 65: [(1,64), (16,49)] }

Checking if an item is in the dictionary is much more efficient than iterating through the list for the same thing. Just look for example if 65 in diccionario .

This makes your comprobar() function unnecessary.

# PRIMERA MEJORA. DICCIONARIOS
def gen_raices1(min, max):
    diccionario = {}

    # Añade a la lista las sumas de raices.
    e1 = 1
    while e1 ** 2 <= max:
        for e2 in range(1, max):
            r1 = e1 ** 2
            r2 = e2 ** 2
            if r1+r2 > max:
              continue
            if r1+r2 not in diccionario:
              diccionario[r1+r2] = []
            if (r2,r1) not in diccionario[r1+r2]:
              diccionario[r1+r2].append((r1,r2))
        e1 += 1
    return diccionario

We want to avoid that repeated values ​​go to the same dictionary key (that is, if we already know that 17 is 1+16 , we do not also enter 16+1 ). That is why I check that (r2, r1) not already in the list associated with diccionario[r1+r2] .

Another possibility is to use sets instead of lists. A set is a type of container that, even though you put repeated data into it, only keeps one copy of each. Therefore we could make the dictionary values ​​are sets and put the ordered tuples there (like this (1,16) and (16,1) , once ordered, both come out (1,16) and the set will only put it once .

The code would look like this (I have my doubts that this is more readable, I almost prefer the previous version):

# SEGUNDA OPCIÓN. DICCIONARIOS y CONJUNTOS
def gen_raices2(min, max):
    diccionario = {}

    # Añade a la lista las sumas de raices.
    e1 = 1
    while e1 ** 2 <= max:
        for e2 in range(1, max):
            r1 = e1 ** 2
            r2 = e2 ** 2
            if r1+r2 > max:
              continue
            if r1+r2 not in diccionario:
              diccionario[r1+r2] = set()
            diccionario[r1+r2].add(tuple(sorted((r1,r2))))
        e1 += 1
    return diccionario

Finally, note that we have to see if r1+r2 was already in the dictionary or not, since if it was not, we must create a new entry for it (a list or an empty set), while if it was already there, what is there to do is add a new pair to the list (or set).

This is simplified by making use of a defaultdict . This is a dictionary that, when you try to access a key that it doesn't have, automatically creates one of the type you want. For example, if we continue with the case where we are using sets:

# TERCERA MEJORA. DEFAULTDICT y CONJUNTOS
from collections import defaultdict

def gen_raices3(min, max):
    diccionario = defaultdict(set)

    # Añade a la lista las sumas de raices.
    e1 = 1
    while e1 ** 2 <= max:
        for e2 in range(1, max):
            r1 = e1 ** 2
            r2 = e2 ** 2
            if r1+r2 > max:
              continue
            diccionario[r1+r2].add(tuple(sorted((r1,r2))))
        e1 += 1
    return diccionario

Notice that we no longer have the if r1+r2 not in diccionario . We directly access the element [r1+r2] as if it existed to update it. If it does not exist, an empty set will automatically be created and then the tuple in question will be added to it.

Reduce unnecessary iterations

In your code you vary r1 from 1 to max , and the same for r2 . This causes you to be computing a lot of unnecessary cases. All those for which r1**2+r2**2 goes over max are discarded, but by thinking a bit about what the indices of the loops should be, you could have avoided calculating them.

Note that it is enough to iterate r1 from 1 to the square root of max , since for every r1 greater than that square root you will have that r1**2 is already greater than max .

On the other hand, it is enough to iterate r2 starting at r1 instead of starting at 1, to avoid generating duplicates like (1,16) and (16,1) . This saves us from having to use sets or having to check if the pair in another order was already saved. It's a great optimization! We can also stop the iteration when reaching the square root of max-r1 , since for any value of r2 greater than that, the sum r1**2+r2**2 would exceed max .

The following code implements these ideas (to calculate the square root I have raised it to 0.5):

# TERCERA MEJORA. ELIMINAR ITERACIONES INNECESARIAS
def gen_raices3(min, max):
    diccionario = defaultdict(list)
    for e1 in range(1, int(max**0.5)):
        for e2 in range(e1, int((max-e1**2)**0.5)+1):
            r1 = e1**2
            r2 = e2 ** 2
            diccionario[r1+r2].append((r1, r2))
    return diccionario

Final optimization

A small extra optimization consists in realizing that the instruction r1 = e1**2 is being executed several times because it is inside the second loop, but since in that second loop r1 does not vary, it always comes out the same, so the we are recalculating unnecessarily. It could be calculated before going into that loop:

# CUARTA MEJORA. LIGERA OPTIMIZACION
from collections import defaultdict

def gen_raices4(min, max):
    diccionario = defaultdict(list)
    for e1 in range(1, int(max**0.5)):
        r1 = e1**2
        for e2 in range(e1, int((max-e1**2)**0.5)+1):
            r2 = e2 ** 2
            diccionario[r1+r2].append((r1, r2))
    return diccionario

Result

In any of the above cases, the result returned by the function is a dictionary (a defaultdict in fact). If you just print it, it will not be ordered from lowest to highest by its keys (since in python the dictionaries do not have a preset order). If you want it to be ordered, you can convert the dictionary into a list of tuples (the first element of the tuple would be the key, the next element its value), and sort that list.

For example:

resultado = gen_raices4(1,100)
print(sorted(resultado.items()))

To get:

[(2, [(1, 1)]), (5, [(1, 4)]), (8, [(4, 4)]), (10, [(1, 9)]), (13, [(4, 9)]), (17, [(1, 16)]), (18, [(9, 9)]), (20, [(4, 16)]), (25, [(9, 16)]), (26, [(1, 25)]), (29, [(4, 25)]), (32, [(16, 16)]), (34, [(9, 25)]), (37, [(1, 36)]), (40, [(4, 36)]), (41, [(16, 25)]), (45, [(9, 36)]), (50, [(1, 49), (25, 25)]), (52, [(16, 36)]), (53, [(4, 49)]), (58, [(9, 49)]), (61, [(25, 36)]), (65, [(1, 64), (16, 49)]), (68, [(4, 64)]), (72, [(36, 36)]), (73, [(9, 64)]), (74, [(25, 49)]), (80, [(16, 64)]), (82, [(1, 81)]), (85, [(4, 81), (36, 49)]), (89, [(25, 64)]), (90, [(9, 81)]), (97, [(16, 81)]), (98, [(49, 49)]), (100, [(36, 64)])]

Execution times

As a curiosity, I have timed how long each of these versions takes to execute, using timeit (which executes the function 1000 times and keeps the average of the best three, to eliminate random "noise"). This is what I got:

  • Your version: 745 µs
  • Version 1 (list dictionaries): 610 µs
  • Version 2 (set dictionaries): 641 µs
  • Version 3 (reduce iterations): 44.7 µs
  • Version 4 (final optimization): 36.7 µs

conclusion

Using more suitable types such as the dictionary reduces the complexity of the code that is easier to read, but it does not reduce the execution time so much, just about 100 µs (in the case of using set() instead of lists, things get slightly worse) .

Instead, thinking about how to eliminate unnecessary iterations improves results dramatically, reducing execution time by an order of magnitude (divide by 10). Note that in this example it is microseconds, and it seems that it is not worth it, but if the problem were longer and your version took 20 minutes to finish, the optimized version would take 1 minute. Compensate!

In any case, as the great sage Donald Knuth said, Premature optimization is the root of all evil . That is, start by making code readable and easy to understand and only if you really need it to go much faster, consider how to change it to improve its speed. Give priority to readability over speed. After all that's why we use python instead of C! 😉

Scroll to Top