Magic Square with python, using itertools

2k views Asked by At

I have some code that will determine if a N*N list of integers forms a magic square:

import itertools

#Function square magic
def magic_square(matrix):
    dimension = len(matrix[0])
    sum_list = []

    #Horizontal code:
    sum_list.extend([sum (lines) for lines in matrix])   

    #Vertical code:
    for col in range(dimension):
        sum_list.append(sum(row[col] for row in matrix))

    #Diagonals code
    diagonal1 = 0
    for i in range(0,dimension):
        diagonal1 +=matrix[i][i]
    sum_list.append(diagonal1)  

    diagonal2 = 0
    for i in range(dimension-1,-1,-1):
        diagonal2 +=matrix[i][i]
    sum_list.append(diagonal2)

    if len(set(sum_list))>1:
        return False
    return True

m=[[7, 12, 1, 14], [2, 13, 8, 11], [16, 3, 10, 5], [9, 6, 15, 4]] 
print(magic_square(m))

m=[[2, 7, 6], [9, 5, 1], [4, 3, 8]]
print(magic_square(m))

m=[[2, 7, 6], [9, 5, 1], [4, 3, 7]]
print(magic_square(m))

print("**************************")

#Now, i use itertools like this:
for i in itertools.combinations(range(1,10), 3):
    if sum(i) == 15:
        print (i)
# I get the combinations each of three numbers with sum 15

My problem is the last part: I want to generate all permutations of the integers 1 through N^2, break each into a square -- a 2-D list of N rows and N columns -- and use my function to find all magic squares. The itertools code I wrote finds combinations of 3 numbers that will do the job, but I can't figure out the combinatorics to form squares.

Thanks @Prune for your help.

If i have:
[1 5 9]
[1 6 8]
[2 4 9]
[2 5 8]
[2 6 7]
[3 4 8]
[3 5 7]
[4 5 6]
how i can generate an square magic and know if it's True o False, using three at a time the elements of matrix?
Example:
[[1 5 9],[1 6 8], [2 4 9]]
or
[[1 5 9],[1 6 8], [2 5 8]]
or
[[1 5 9],[1 6 8], [2 6 9]] Etcetera, etcetera.

3

There are 3 answers

5
Prune On

I see -- you want to generate all permutations of magic squares. You need to cover all permutations of the range from 1 through N^2, returned as N lists of N elements each.

import itertools

N = 3
for seq in itertools.permutations(range(1, N*N+1)):
    # Split the sequence into a candidate magic square,
    #   N rows of N elements each.
    cand = [seq[i:i+N] for i in range(0, N*N, N)]

This produces a series of candidate squares; feed each one, in turn, to your checking routine, and print the ones that come up True. I expect that you can handle that part.

Here are a few sample candidates from the early part of the generation:

[(1, 3, 5), (6, 2, 8), (4, 7, 9)]
[(1, 3, 5), (6, 2, 8), (4, 9, 7)]
[(1, 3, 5), (6, 2, 8), (7, 4, 9)]
[(1, 3, 5), (6, 2, 8), (7, 9, 4)]
[(1, 3, 5), (6, 2, 8), (9, 4, 7)]
[(1, 3, 5), (6, 2, 8), (9, 7, 4)]
[(1, 3, 5), (6, 2, 9), (4, 7, 8)]
[(1, 3, 5), (6, 2, 9), (4, 8, 7)]
[(1, 3, 5), (6, 2, 9), (7, 4, 8)]
[(1, 3, 5), (6, 2, 9), (7, 8, 4)]
[(1, 3, 5), (6, 2, 9), (8, 4, 7)]
[(1, 3, 5), (6, 2, 9), (8, 7, 4)]

Change of method

Note that this is not your original algorithm: this generates whole squares, rather than just rows of 3. The independent generation has a logic flaw, in that it will generate magic squares that do not include all 9 numbers, while duplicating others. For instance:

7 2 6
4 5 6
4 8 3
0
Prune On

Algorithm from given stopping point

You currently have all eight possible combinations of three distinct integers 1-9 that sum to 15. To solve the magic square in the straightforward way you requested, I suggest these steps:

  • Generate all permutations of each combination: six per combination, for a total of 48. One way to do this is to generate all permutations (instead of combinations) in your current code.
  • Consider these rows in permutations of 3 at a time. For each permutation, check to see whether stacking the rows in the generated order forms a square. You can check:
    • Does each of the three columns sum to 15?
    • Does each diagonal sum to 15?
    • Are the 9 numbers unique?

Faster Code

There are various ways to attack the permutations for efficiency. For instance, divide the rows into four groups by lowest element (1, 2, 3, or 4). In generating your squares, pick no more than one row from each group. This will greatly reduce the total squares you check, since it reduces duplication of elements.

Another way is to pick the first two rows, and then derive the third row from column sums. Then you have only four checks to make: that the two diagonals sum to 15, that the generated bottom row is legal (has only numbers 1-9), and that there are no duplicated numbers.

Finding the 8 rows more efficiently

You don't have to paw through the 720 triples to find the 8 rows. Instead, generate the 90 starting pairs; for each, derive the third element (15 minus the first two). If the third element is one of the 7 missing numbers (1-9, but neither of the first two), then it's one of your desired rows.

I hope that this leads you to a solution.

0
FoConrad On

Here is a solution in a quasi one-liner (with the benefit of producing answers as they are found). It maps the permutations of length N*N to numpy reshape, then determines if the matrix is magic.

import numpy
import functools
import itertools

N = 3
for item in filter(lambda o: len(set(numpy.sum(o, axis=0))
                                 .union(numpy.sum(o, axis=1))
                                 .union({o.diagonal().sum()})
                                 .union({numpy.fliplr(o).diagonal().sum()})
                                 ) == 1,
                   map(functools.partial(numpy.reshape, newshape=(N, N)), 
                       itertools.permutations(range(N*N)))):
    print(item)