Solving Magic Squares via Backtracking

This post is part of a series.

Brute force search can often be slow or infeasible when there are lots of possibilities that must be checked.

However, a technique called backtracking can often be used to drastically cut down the number of possibilities that must be checked. The idea is that whenever we find ourselves constructing a set of possibilities that we know are hopeless, we skip over them and backtrack to the the point right before we started constructing the hopeless possibilities.

Exercise: Brute Force

Use brute-force search to find all arrangements of the digits $1,$ $2,$ $\ldots,$ $9$ into a $3 \times 3$ magic square where all the rows, columns, and diagonals add up to $15$ and no digits are repeated.

You may use $9$ nested for loops if you’d like. It’s ugly but conceptually simple and gets the job done. To illustrate, pseudocode is provided below.

digits = [1, 2, ..., 9]
square = [
    [None, None, None]
    [None, None, None]
    [None, None, None]
]

for num1 in digits:
  clear out the square and put num1 in

  for num2 in digits excluding num1:
    clear out the square and put num1, num2 in

    for num3 in digits excluding num1, num2:
      clear out the square and put num1, num2, num3 in

      ... and so on

      for num9 in digits excluding num1, ..., num8:
        clear out the square and put num1, ..., num9 in

        if is_valid(square):
          print(square)

Note that this solution will be rather slow because it will require checking $9! = 362\,880$ arrangements of digits.

Exercise: Backtracking

Repeat the above exercise, but this time, whenever you reach an arrangement of digits that can no longer become a valid magic square, do not explore that arrangement any further. You can accomplish this by doing the following:

1. Write a function is_hopeless(square) to check whether an incomplete square could ever become a valid square. A square is hopeless if any row, column, or diagonal is filled in and does not sum to $15.$
2. Place continue statements in your $9$ nested for loops so that you skip inner loops whenever you detect a hopeless square.

Below are some examples to illustrate how the is_hopeless function should work.

>>> is_hopeless([
        [1,    2,    None],
        [None, 3,    None],
        [5,    6,    4   ]
    ])

OUTPUT: True 

REASONING
- A diagonal is filled in and it doesn't sum to 15.

- No matter how we fill the rest of the square,
  this diagonal still won't sum to 15.

- Therefore, this incomplete arrangement cannot lead
  to any valid square.

>>> is_hopeless([
        [1,    2,    None],
        [None, None, None],
        [5,    6,    4   ]
    ])

OUTPUT: False 

REASONING: There are no filled rows, columns, or
diagonals that don't sum to 15.

>>> is_hopeless([
        [None, None, None]
        [None, None, None]
        [None, None, None]
    ])

OUTPUT: False 

REASONING: There are no filled rows, columns, or
diagonals that don't sum to 15.

Below is an illustration of how to skip inner loops using continue statements.

for num1 in digits:
    clear out the square and put num1 in it

    if is_hopeless(square):
        continue

    for num2 in digits excluding num1:
        clear out the square and put num1, num2 in it

        if is_hopeless(square):
            continue

        ... and so on

To give a concrete demonstration of backtracking, the first several arrangements that get explored are shown below.

[[_,_,_],
 [_,_,_],
 [_,_,_]]

[[1,_,_],
 [_,_,_],
 [_,_,_]]

[[1,2,_],
 [_,_,_],
 [_,_,_]]

[[1,2,3],
 [_,_,_],
 [_,_,_]]
^ hopeless -- do not explore further

[[1,2,4],
 [_,_,_],
 [_,_,_]]

^ hopeless -- do not explore further

[[1,2,5],
 [_,_,_],
 [_,_,_]]

^ hopeless -- do not explore further

...

[[1,2,9],
 [_,_,_],
 [_,_,_]]

^ hopeless -- do not explore further

[[1,3,_],
 [_,_,_],
 [_,_,_]]

...

This post is part of a series.