python magic square finder for arbitrary numbers

Question

I wrote this code to get all possible magic squares that can be made with a list you put in. The only contraints are that the count of numbers is a square and that each number is unique. It even works for complex numbers.

I dont know how good the performance is, it needs about 1.5 to 2hours to find all 7040 4*4 magic squares with numbers from 1 to 16.

The finding strategy is to get all possible permutations of all subsets of the number set that sum to the magic number and then put these as a whole row into the matrix, recursively.

All feedback is welcome :)

from copy import deepcopy
from itertools import combinations, permutations
from math import sqrt
def main(numbers):
 '''Finds all magic squares that can be made with
 the given set of numbers.'''
 global dim, magicnumber, emptyrow
 dim = int(sqrt(len(numbers)))
 magicnumber = sum(numbers) / dim
 emptyrow = ["" for _ in range(dim)]
 current = [emptyrow for _ in range(dim)]
 groups = possibilities(numbers, dim, magicnumber)
 placeRow(current, groups, row=0)
 report(solutions)
def possibilities(numbers, dim, magicnumber):
 '''Returns a list of all the ways to reach
 the magic number with the given list of numbers.'''
 combos = [
 list(x) for x in list(
 combinations(
 numbers,
 dim)) if sum(x) == magicnumber]
 possibilities = [list(permutations(x)) for x in combos]
 possibilities = [[list(y) for y in x] for x in possibilities]
 possibilities = [item for sublist in possibilities for item in sublist]
 return possibilities
def remainding_possibilities(matrix, possibilities):
 '''Returns the remainding possibilities once the matrix has entries.'''
 entries = {item for sublist in matrix for item in sublist}
 remainders = [x for x in possibilities if entries.isdisjoint(x)]
 return remainders
def placeRow(matrix, groups, row=0):
 '''Recursive function that fills the matrix row wise
 and puts magic squares into "solutions" list.'''
 godeeper = False
 current = matrix
 for group in groups:
 current[row] = group # put the whole group into the row
 if emptyrow in current:
 remainders = remainding_possibilities(current, groups)
 godeeper = placeRow(current, remainders, row=row + 1)
 else:
 if check(current):
 solutions.append(deepcopy(current))
 current[row] = emptyrow
 return False
 else:
 current[row] = emptyrow
 if godeeper is False:
 current[row] = emptyrow
 return False
def check(matrix):
 '''Returns false if current matrix is not or cant be made
 into a magic square.'''
 # rows
 # not needed because we fill row wise
 # for x in range(dim):
 # if "" not in matrix[x]:
 # if sum(matrix[x]) != magicnumber:
 # return False
 # only if we have positive numbers only
 # else:
 # if sum(transposed[x]) > magicnumber:
 # return False
 # diagonals
 diag1 = [matrix[x][x] for x in range(dim)]
 if "" not in diag1:
 if sum(diag1) != magicnumber:
 return False
 # only if we have positive numbers only
 else:
 if sum(diag1) > magicnumber:
 return False
 diag2 = [matrix[x][dim - 1 - x] for x in range(dim)]
 if "" not in diag2:
 if sum(diag2) != magicnumber:
 return False
 # only if we have positive numbers only
 else:
 if sum(diag2) > magicnumber:
 return False
 # columns
 transposed = transpose(matrix)
 for x in range(dim):
 if "" not in transposed[x]:
 if sum(transposed[x]) != magicnumber:
 return False
 # only if we have positive numbers only
 else:
 if sum(transposed[x]) > magicnumber:
 return False
 return True
def transpose(matrix):
 '''Transpose a matrix.'''
 return list(map(list, zip(*matrix)))
def report(solutions):
 ''' Writes solutions to text file.'''
 with open(f"solutions.txt", 'w') as txtfile:
 txtfile.write(
 f"Found {len(solutions)} magic squares:\n\n")
 for solution in solutions:
 for line in solution:
 for entry in line:
 txtfile.write(f"{entry}" + " ")
 txtfile.write("\n")
 txtfile.write("\n")
if __name__ == "__main__":
 # Some inputs for main().
 complex3 = [(complex(x, y))
 for x in range(1, 4)
 for y in range(1, 4)]
 complex4 = [(complex(x, y))
 for x in range(1, 5)
 for y in range(1, 5)]
 complex5 = [(complex(x, y))
 for x in range(1, 6)
 for y in range(1, 6)]
 test3 = [x for x in range(1, 10)]
 test4 = [x for x in range(1, 17)]
 test5 = [x for x in range(1, 26)]
 solutions = []
 main(complex3)

Answer 1

Just reviewing the function possibilities.

This function is doing a lot of unnecessary extra work and I think it would be worth your while trying to understand why you made it so complicated.

The first steps are: (i) find all combinations of dim elements of numbers; (ii) convert the combinations to a list; (iii) filter for combinations that add up to the magic number; (iv) convert each combination to a list.

combos = [
 list(x) for x in list(
 combinations(
 numbers,
 dim)) if sum(x) == magicnumber]

What is the purpose of steps (ii) and (iv)? If you omit these steps then this line simplifies to:

combos = [x for x in combinations(numbers, dim) if sum(x) == magicnumber]

and this works just as well as the original.

The next steps are: (v) find the permutations of each combination; (vi) turn each group of permutations into a list; (vii) turn each permutation into a list; (viii) flatten the list of lists of permutations into a single list.

possibilities = [list(permutations(x)) for x in combos]
possibilities = [[list(y) for y in x] for x in possibilities]
possibilities = [item for sublist in possibilities for item in sublist]

But again, steps (vi) and (vii) are unnecessary, and steps (v) and (viii) can be combined into one, leaving you with:

possibilities = [p for x in combos for p in permutations(x)]

and the definition of combos can be inlined here, getting:

return [p for c in combinations(numbers, dim)
 if sum(c) == magicnumber
 for p in permutations(c)]

• Tweakimp: Thank you for your answer. The code just "grew" that way and then I couldnt refactor it to a shorter version like you did. Would you be so kind to also look at remainding_possibilities, because most of the run time comes from that function.
• Gareth Rees: When code "just grows" that means that you don't fully understand how it works and you've lost control of it. The way to regain understanding is to work through the steps in detail on a small example, so that you can see how each step transforms the data structures.
• Tweakimp: I dont think we are talking about the same thing. What I meant was as I was writing it, I changed it many times until it worked and then I couldnt go back or rephrase it because I am not good enough to say it in another way. That is exactly why I post here.