Program to find magic square of squares

I have a program which finds, well, a magic square of squares. The problem is, it's quite slow - it processes around 50 numbers in half a second when the number range is above 40,000.

Is there any way to improve this code?

```
from math import floor

# checksquare checks for the possibility of a set being a magic square.
# One example of an incoming list (split for readability) is:
# [[1634],
# [15, 25, 28], [11, 27, 28], [8, 27, 29],
# [3, 28, 29], [12, 23, 31], [13, 21, 32],
# [9, 23, 32], [4, 23, 33], [3, 20, 35]]
# The square of each number in the lists of length 3
# add up to the single number in the list of length 1.
# (e.g. 15^2 + 25^2 + 28^2 = 1634)
# This function checks for this possibility by doing so:
# If all of the numbers in the sets of 3 are repeated at least once,
# Then it outputs it to a separate file.
def checksquare(listin):
# listcheck only contains the lists of length 3.
listcheck = listin[1:]
# dictofnos is used to check the amount of each number in listcheck.
dictofnos = {}
# setof0s is used to remove numbers which are not repeated.
setof0s = []
# The first loop checks the amount of each number in listcheck.
for e in range(len(listcheck)):
for f in range(3):
try:
dictofnos[(listcheck[e])[f]] += 1
except KeyError:
dictofnos[(listcheck[e])[f]] = 1
# The second loop removes any value that is not repeated.
for g in dictofnos:
if dictofnos[g] == 1:
for h in range(len(listcheck)):
if g in listcheck[h]:
for j in range(3):
if dictofnos[(listcheck[h])[j]] == 0:
pass
dictofnos[(listcheck[h])[j]] -= 1
listcheck.remove(listcheck[h])
listcheck.append([])
for i in range(len(listcheck)):
if len(listcheck[i]) != 3:
listcheck.remove(listc

This is at least a start. I did not look very heavily in finding a better algorithm itself, just making small improvements of you algorithm. Nevertheless this resulted in a speed-up of about 30%.

For comparison (I changed the call to powers(50) so it does not take so long):

# Your code
$ python -m cProfile magic_square_orig2.py
         81524 function calls in 3.462 seconds
# With the changes below
$ python -m cProfile magic_square2.py     
         18110 function calls (18107 primitive calls) in 2.496 seconds

With powers(75):

# Your code
$ python -m cProfile magic_square_orig2.py
         336810 function calls in 30.210 seconds
# With the changes below
$ python -m cProfile magic_square2.py
         148883 function calls (148880 primitive calls) in 18.832 seconds

Use better names (!!)

Even now, where you changed some variable names for better names, the function checksquare is very hard to understand because of variables called g, h, j, k. But since I will take away their meaning as integers below, anyways, we can find better names for them.

Iterate over the contents of a list

It is always easier to iterate over the list, instead of iterating an index, compare:

for i in range(len(l)):
    print l[i]
for element in l:
    print element

The latter is a lot easier to read (and understand). It is the recommended way to iterate over a list. So I changed your logic and gave the variables better names:

Use collections.Counter()

There is already an existing construct that builds a dictionary with counts of objects, it is callen collections.Counter. This way you can replace:

for e in range(len(listcheck)):
        for f in range(3):
            try:
                dictofnos[(listcheck[e])[f]] += 1
            except KeyError:
                dictofnos[(listcheck[e])[f]] = 1

with:

from collections import Counter
...
dictofnos = Counter(item for sublist in listcheck for item in sublist)

Use list comprehensions where possible

List comprehensions are in general faster than manually writing the for loop, because they are implemented in C (they are not, see e.g. here) a more succinct way to write simple loops that build a list. You can replace some loops:

for i in range(len(listcheck)):
                if len(listcheck[i]) != 3:
                    listcheck.remove(listcheck[i])

becomes one of these two:

listcheck = filter(lambda powers: len(powers) == 3, listcheck)
listcheck = [powers for powers in listcheck if len(powers) == 3]

which seem to be about the same speed. where the latter should be slightly faster, because the filter has to take a lambda instead of a predefined function. But using the fact that the only two possible lengths are 0 and 3 and bool(0) == False and bool(3) == True we can just use

listcheck = filter(len, listcheck)

in this case.

function powers

You compute for example b2 more than once. It saves quite some time if you save b2 = b2 (and similar for the other variables) at appropriate places.

int() already performs floor so it is not needed here. (int(3.14) == 3 and int(3.99) == 3)

You can collect what to write to the output file and write it in one go. This should be faster than repeated opening, writing and closing of the file. For this the function checksquare needs to be adapted to return the values instead of writing it:

if listcheck:
        return listin[0], dictofnos, listcheck

and in powers we add a list to collect the return values:

def powers(limit):
    out = []
    ....
        if temp >= 8:
                squares = checksquare(templist)
                if squares:
                    out.append(squares)
    ....
    return out

Additionally, we can put the writing part to a new function, separating the calculation and output part:

def write_powers(n):
    with open("output2.txt", "w") as out_file:
        for power in powers(n):
            out_file.write("\n{}\n{}\n{}\n".format(*power))

This also guarantees that the file will be overwritten with each subsequent call to the script ("w" writes, over-writing the file, while you "a+" was always appending).

Misc

Use the __name__ hook in order to allow importing you function from another script without executing powers(75) every time:

if __name__ == "__main__":
    write_powers(75)

Use an actual set for setof0s.

Result

```
from collections import Counter

def checksquare(listin):
# listcheck only contains the lists of length 3.
listcheck = listin[1:]
# dictofnos is used to check the amount of each number in listcheck.
dictofnos = Counter(factor for factors in listcheck for factor in factors)
# setof0s is used to remove numbers which are not repeated.
setof0s = set()
# The first loop checks the amount of each number in listcheck.
# The second loop removes any value that is not repeated.
for factor in dictofnos:
if dictof

StackExchange Code Review Q#135260, answer score: 7