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(listcheck[i])
    # This if/elif is used to catch any lists that passed the two loops
    # while having non-repeating numbers.
    if 0 in dictofnos.values():
        [setof0s.append(k) for k in dictofnos if dictofnos[k] == 0]
    elif 1 in dictofnos.values():
        listcheck = listin[:1] + listcheck
        checksquare(listcheck)
        return None
    # The final loop is deleting entries in the dict which are removed
    # (hence the use of setof0s).
    for l in setof0s:
        del dictofnos[l]
    # Outputs to output.txt.
    if len(listcheck) != 0:
        output = open("output.txt", "a+")
        output.write("\n" + str(listin[0]) + "\n" + str(dictofnos) + "\n" + str(listcheck) + "\n")
        output.close()


# powers checks if 3 squares add up to a number.
def powers(limit):
    numberrange = 3 * ((limit + 1) ** 2)
    for a in range(numberrange):
        # Everything here is for efficiency.
        temp = 0
        templist = [[a]]
        check = 1
        b = 0
        c = int(floor((a / 4.0) ** 0.5))
        d = int(floor((a / 2.0) ** 0.5))

        while b &lt;= c &lt;= d &lt;= limit and d ** 2 &lt; a:
            b += check
            if b == c or (b * 2) ** 2 &gt; a:
                c += 1
                b = 1
                continue
            if c == d:
                d += 1
                c = int(floor((a / 4.0) ** 0.5))
                b = 1
                continue
            if (b + c + d) % 2 != a % 2:
                check = 2
                b -= 1
                continue
            if b ** 2 + c ** 2 + d ** 2 == a:
                templist.append([b, c, d])
                temp += 1
        # If 8 solutions for b^2 + c^2 + d^2 = a are found
        # for any a, then it is sent to checksquare.
        if temp &gt;= 8:
            checksquare(templist)

# The usage of these functions would be to put
# powers(whatever the upper bound of d is).
powers(100)


EDIT: Added explanation (sorry if it's really long), and improved some minor things.
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 b**2 more than once. It saves quite some time if you save b2 = b**2 (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 &gt;= 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 dictofnos[factor] == 1:
            for factors in listcheck:
                if factor in factors:
                    for power in factors:
                        if dictofnos[power] == 0:
                            pass
                        dictofnos[power] -= 1
                    listcheck.remove(factors)
            #listcheck = filter(lambda powers: len(powers) == 3, listcheck)
            listcheck = [powers for powers in listcheck if len(powers) == 3]
    # This if/elif is used to catch any lists that passed the two loops
    # while having non-repeating numbers.
    if 0 in dictofnos.values():
        setof0s |= set(k for k in dictofnos if dictofnos[k] == 0)
    elif 1 in dictofnos.values():
        listcheck = listin[:1] + listcheck
        checksquare(listcheck)
        return
    # The final loop is deleting entries in the dict which are removed
    # (hence the use of setof0s).
    for l in setof0s:
        del dictofnos[l]
    # Outputs to output.txt.
    if listcheck:
        return listin[0], dict(dictofnos), listcheck


# powers checks if 3 squares add up to a number.
def powers(limit):
    out = []
    numberrange = 3 * ((limit + 1) ** 2)
    for a in range(numberrange):
        # Everything here is for efficiency.
        temp = 0
        templist = [[a]]
        check = 1
        b = b2 = 0
        c = int((a / 4.0) ** 0.5)
        d = int((a / 2.0) ** 0.5)
        c2, d2 = c**2, d**2

        while b &lt;= c &lt;= d &lt;= limit and d2 &lt; a:
            b += check
            b2 = b**2
            if b == c or 4*b2 &gt; a:
                c += 1
                c2 = c**2
                b = b2 = 1
                continue
            if c == d:
                d += 1
                d2 = d**2
                c = int((a / 4.0) ** 0.5)
                c2 = c**2
                b = b2 = 1
                continue
            if (b + c + d) % 2 != a % 2:
                check = 2
                b -= 1
                b2 = 1
                continue
            if b2 + c2 + d2 == a:
                templist.append([b, c, d])
                temp += 1
        # If 8 solutions for b^2 + c^2 + d^2 = a are found
        # for any a, then it is sent to checksquare.
        if temp &gt;= 8:
            squares = checksquare(templist)
            if squares:
                out.append(squares)
    return out

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

# The usage of these functions would be to put
# powers(whatever the upper bound of d is).
if __name__ == "__main__":
    write_powers(75)