Optimize "Fill magic square" in python with backtracking

I have written a python program with backtracking that fills a n * n magic square. It solves for n = 4 in 4.5 seconds but gets stuck when I run it for n = 5 on my machine. How can I optimize the algorithm to make it run faster?

from time import perf_counter

def fillMatrix(n):
nSquared = n * n
lineSum = n * (nSquared + 1) / 2

candidates = set(range(1, nSquared + 1))
matrix = [[None for _ in range(n)] for _ in range(n)]
def isValid(row, col):
# row
rowSum = 0
isFull = True
for item in matrix[row]:
if not item:
isFull = False
continue
rowSum += item
if rowSum > lineSum:
return False
if isFull and rowSum != lineSum:
return False

# column
colSum = 0
isFull = True
for i in range(n):
item = matrix[i][col]
if not item:
isFull = False
continue
colSum += item
if colSum > lineSum:
return False
if isFull and colSum != lineSum:
return False

# diagonal
if row != col and row + col != n - 1:
return True

diagSum = 0
isFull = True
for i in range(n):
item = matrix[i][i]
if not item:
isFull = False
continue
diagSum += item
if diagSum > lineSum:
return False
if isFull and diagSum != lineSum:
return False

diagSum = 0
isFull = True
for i in range(n):
item = matrix[n - i - 1][i]
if not item:
isFull = False
continue
diagSum += item
if diagSum > lineSum:
return False
if isFull and diagSum != lineSum:
return False

return True

def solve(row, col):
if matrix[row][col] == None:
for candidate in candidates.copy():
matrix[row][col] = candidate
if not isValid(row, col):
matrix[row][col] = None
continue
candidates.remove(candidate)

if row == n - 1 and col == n - 1:
return True

currentSolution = False
if col == n - 1:
currentSolution = solve(row + 1, 0)
else:
currentSolution = solve(row, col + 1)
if currentSolution:
return True
matrix[row][col] = None
return False
return True

t1 = perf_counter()
print(solve(0, 0))
print(matrix)
t2 = perf_counter()
print(f'{t2 - t1}s')

fillMatrix(5)


Optimization is pointless

Random magic square hunting peters out very quickly. If n is the order of a magic square (the number of cells in a row), the domain to be searched for valid magic squares grows at the rate of (n*n)!. If we conceptualize a magic square as a one-dimensional list of length n*n, a naive brute-force search would iterate over all permutations of that list. For n=3 that is doable, because there are only 362880 lists to check. But for n=4, that domain becomes about 21 trillion. Since there are known to be 880 4x4 magic squares, that means random hunting in that space will emit a hit once per 23.8 billion checks. Python cannot do that.

Even smart searches don't survive long in that environment. Your implementation takes some advantage of the known constraints of the problem to short-circuit the searching and speed things up. But there are pressing limits here as well. Consider an algorithm that leverages the constraints even more effectively than your code. Instead of filling in the grid one cell and at time and then checking for violated constraints, one could pre-compute all possible ways of creating a valid row/column/diagonal, given the size of the magic square (n), the implied magic constant (the needed sum for every row/column/diagonal), and its universe of eligible numbers (1..n inclusive). One could also pre-compute a lookup index mapping any partially-completed row/column/diagonal to the other numbers that would complete it in a valid way. In addition, the algorithm could proceed in a way that maximizes the interactions among the constraints: first fill the diagonals, then row 0, then column 0, row 1, column 1, etc. The benefit of that kind of crossing approach is that each prior step imposes additional constraints on subsequent steps, thus reducing the number of viable sums that have to be searched. Such an algorithm would be considerably faster for various reasons: (1) it would fill in valid rows/columns/diagonals in a single shot (rather than one cell at a time); (2) no checking for validity would be required because the algorithm would be premised on only filling in valid rows/columns/diagonals; (3) most important, the scope of the search space would be much smaller since we would only be considering combinations of numbers that achieve valid sums rather than naively filling in the next cell and then checking for violations. Unfortunately, none of that is enough.

# Let's explore magic squares of size 3 through 7 to see their
# magic constant along with the number of valid ways that a
# row/column/diagonal can add up to that constant.

from itertools import combinations

for n in range(3, 8):
n_stop = n * n + 1
magic_constant = int(n * n_stop / 2)
universe = tuple(range(1, n_stop))
valid_rows = tuple(
c
for c in combinations(universe, n)
if sum(c) == magic_constant
)
print(n, magic_constant, len(valid_rows))

# Output.

3 15 8
4 34 86
5 65 1394
6 111 32134
7 175 957332


By the time we get n=7 even our smart algorithm will be utterly swamped. For magic squares of size 7 there are nearly 1M ways to get the needed sum of 175. And even that number is a big underestimate of the size of the domain. Because that code snippet uses combinations() it is normalizing the valid ways to make the needed sums. When filling in any particular row/column/diagonal, we would actually need to check every permutation of the current combination being examined. Furthermore, the magnitudes printed above represent only the size of the domain at the top level (where we fill in the first diagonal). That magnitude would need to be multiplied by the constraint-surviving permutations that occur at deeper levels as we fill in the grid according to our crossing plan. As a result, I'm pretty confident that our envisioned algorithm could handle n=5, but I am quite pessimistic about n=6, and n=7 seems impossible.

There are a variety of square-generating algorithms. A different approach is to grab one of the known algorithms to generate specific magic squares -- for example, up and to the right. And if that seems too boring, one could also take a square generated in such a manner and transform it in a variety of ways (also see).

• Lots of great insights in your answer! I know there are other methods for magic square generation but I am only practising my problem-solving skill for coding interviews and was wondering if there's a way to improve the backtracking performance(which you did answer in 2nd paragraph.) Jan 9 at 5:13
• @shawon191 Good luck with your practice. It was an interesting puzzle to learn more about. I started exploring it with much greater optimism for finding a clever ways achieve a slightly higher n, and it was interesting to be reminded how prodigious an (n*n)! growth rate really is.
– FMc
Jan 9 at 17:01