Program to find magic square of squares

Question

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 <= c <= d <= limit and d ** 2 < a:
            b += check
            if b == c or (b * 2) ** 2 > 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 >= 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.

Answer 1

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 >= 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 <= c <= d <= limit and d2 < a:
            b += check
            b2 = b**2
            if b == c or 4*b2 > 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 >= 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)

Answer 2

You have nondeterminism in your code, which will make it very difficult to have fair benchmarks. Case in point, CPython produces different results to PyPy.

We can fix this by removing said nondeterminism. I believe your only significant source is

for g in dictofnos:

since iteration order is arbitrary. I changed this to

for g in sorted(dictofnos, reverse=True):

I'm using reverse since it seems to give more outputs than the other way around, which I assume is preferred.

Now we can get useful timings. The first thing to try is how fast PyPy is:

$$ \begin{array}{ll} \text{runtime} & \text{time}/s \\ \hline \text{CPython} & 3.1\\ \text{PyPy} & 110 \\ \end{array} $$

Well, it seems we're done making it faster. Let's fix some other things instead.

Firstly, your code is mixing IO with logic. We can fix this by yielding results out to the caller. This allows us to keep IO happening as soon as we already do, but let the caller do it. In checksquare, return the needed values (there are two places). In powers, write

out = checksquare(templist)
if out is not None:
    yield out

Then let the caller do the writing. You don't want the function to do the writing to call powers directly, because that recouples IO and logic (albeit to a much lesser extent) - just pass powers into the function.

Further, just print the output. When running the program, it's trivial to redirect to a file. Thus, printing it makes a lot more sense than hardcoding the file you're printing to. I recommend from __future__ import print_function at this point as well.

Lots of your documentation is in comments above functions rather than docstrings. You should write docstrings instead, and use clearer, more declarative language.

Before we fix checksquare's documentation, though, notice that it's really taking two parameters:

checksquare([[first_argument], some, more, lists])

This would be a lot clearer as

checksquare(first_argument, [some, more, lists])

So I replaced it with

def checksquare(squaresum, triplets):

But this just makes it obvious that squaresum is never used by the function, so shouldn't be passed in in the first place. This means that the comment for the function is also a lie!

So what is checksquare doing? Idunno, let's work it out.

dictofnos tells me nothing about its purpose, but I can see that it's a mapping from numbers inside the triplets to counts of some kind. This suggests a Counter, though that doesn't clarify a lot. The comment doesn't really help either. I'll call it counts for now.

Note that (triplets[e])[f] is just triplets[e][f], but the first loop can just be

for triplet in triplets:
    counts.update(triplet)

Personally I'd just initialize it with

counts = Counter(chain.from_iterable(triplets))

and skip that step, though.

I'm going to change our deterministic iteration to

for val, count in counts.most_common():

because this produces even more outputs and ends up easier. val is another problematic name, but it's better than g.

The if count == 1 check should become

if count != 1:
    continue

Then

if counts[(triplets[h])[j]] == 0:
    pass

does nothing, so should be removed. It's also impossible for it to happen.

Now I tidied up the next loop:

# Remove any value that is NOT repeated
for val, count in counts.most_common():
    if count != 1:
        continue

    for triplet_idx in range(len(triplets)):
        if val not in triplets[triplet_idx]:
            continue

        for to_remove in triplets[triplet_idx]:
            counts[to_remove] -= 1
        triplets.remove(triplets[triplet_idx])
        triplets.append([])

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

It's really odd that you're writing

triplets.append([])

and immidiately trying to remove them. Luckily you never seem to do this twice, because if you did your removal would have shifted off the index that you were checking next! But it's still safer and faster to write

triplets = filter(bool, triplets)

Note that bool is sufficient because the only list lengths are 0 or 3.

There's another remove that causes the shift problem:

triplets.remove(triplets[triplet_idx])
triplets.append([])

Since triplet_idx goes up by one each time and you shift the value down, you end up skipping this one! Try

triplets[triplet_idx] = []

instead. In fact, by this point we should just create a new list in one step:

for val, _ in counts.most_common():
    if counts[val] != 1:
        continue

    def new_triplets():
        for triplet in triplets:
            if val in triplet:
                for to_remove in triplet:
                    counts[to_remove] -= 1
            else:
                yield triplet

    triplets[:] = new_triplets()

But this still isn't great as it takes \$\mathcal{O}(n^2)\$ time - we're iterating over every value, and for each of those over every triplet. Rather try the other way around:

def new_triplets():
    for triplet in triplets:
        if any(counts[val] == 1 for val in triplet):
            for val in triplet:
                counts[val] -= 1
        else:
            yield triplet

triplets[:] = new_triplets()

Note that this doesn't actually correspond to an order done the previous way, since triplets are discarded in order, not by an arbitrary sort of their elements.

Then there is:

if 0 in counts.values():
    [setof0s.append(k) for k in counts if counts[k] == 0]
elif 1 in counts.values():
    return checksquare(triplets)

# The final loop is deleting entries in the dict which are removed
# (hence the use of setof0s).
for l in setof0s:
    del counts[l]

This should really be:

if 0 in counts.values():
    # Remove elements with count of zero
    counts += Counter()
elif 1 in counts.values():
    return checksquare(triplets)

This makes no sense to me, though - why is 1 in counts.values() allowed when 0 in counts.values()?

You can avoid recursion by using a loop:

while True:
    counts = Counter(chain.from_iterable(triplets))

    # Remove any value that is NOT repeated
    def new_triplets():
        ...

    triplets[:] = new_triplets()

    if 0 in counts.values():
        # Remove zero elements in counts
        counts += Counter()
        break

    if 1 not in counts.values():
        break

And you can move counts = outside of the loop.

counts = Counter(chain.from_iterable(triplets))

while 1 in counts.values():
    # Remove any value that is NOT repeated
    def new_triplets():
        ...

    triplets[:] = new_triplets()

    if 0 in counts.values():
        # Remove zero elements in counts
        counts += Counter()
        break

Finally, for this function

if len(triplets) != 0:
    return counts, triplets

is just

if triplets:
    return counts, triplets

but I'd just remove the if triplets because the caller can do that and it makes the return type uniform.

On to powers. We have numberrange, which should be squaresum_limit or inlined. floor is unneeded. Instead of a / 4.0 you can from __future__ import division and do a / 4. Use tuples instead of lists for triplets. Remove temp; it buys you nothing.

I'd like to give more advice about powers but I really don't get what its search strategy is. A lot of the inner loop doesn't make a lot of sense to me. If I were to search for triples such that a = b ** 2 + c ** 2 + d ** 2 and b <= c <= d, I'd do

def limit(n, a, b=0, c=0):
    space = a - b**2 - c**2
    return int(space ** 0.5 / n)

def get_triplets():
    for b in range(limit(3, a) + 1):
        for c in range(b, limit(2, a, b) + 1):
            d = limit(1, a, b, c)

            if a == b ** 2 + c ** 2 + d ** 2:
                yield a, b, c

triplets = list(get_triplets())

but this produces all possible such triples, whereas yours only produces a small subset that I don't really understand, in an order I don't really understand.

So this is what I have. I'm not sure it'll match what you were aiming for, though, as I don't really get what your code does.

from __future__ import division, print_function

from collections import Counter
from itertools import chain


def checksquare(triplets):
    counts = Counter(chain.from_iterable(triplets))

    while 1 in counts.values():
        # Remove any value that is NOT repeated
        def new_triplets():
            for triplet in triplets:
                if any(counts[val] == 1 for val in triplet):
                    for val in triplet:
                        counts[val] -= 1
                else:
                    yield triplet

        triplets[:] = new_triplets()

        if 0 in counts.values():
            # Remove zero elements in counts
            counts += Counter()
            break

    return dict(counts), triplets


def powers(max_squaresum):
    def limit(n, a, b=0, c=0):
        space = a - b**2 - c**2
        return int(space ** 0.5 / n)

    for a in range(3 * ((max_squaresum + 1) ** 2)):
        def get_triplets():
            for b in range(limit(3, a) + 1):
                for c in range(b, limit(2, a, b) + 1):
                    d = limit(1, a, b, c)

                    if a == b ** 2 + c ** 2 + d ** 2:
                        yield a, b, c

        triplets = list(get_triplets())

        if len(triplets) >= 8:
            counts, triplets = checksquare(triplets)
            if triplets:
                yield a, counts, triplets

def display_output(values):
    for squaresum, counts, triplets in values:
        print()
        print([squaresum])
        print(dict(counts))
        print(triplets)

display_output(powers(100))