我的代码可以更简单吗？

Question

An n-by-n square matrix (table of numbers) is a magic matrix if the sum of its row and the sum of each column are identical. For example, the 4-by-4 matrix below is a magic matrix. The sum of every row and the sum of every column are exactly the same value 34.

        16 2 3 13
        5 11 10 8
        9 7 6 12
        4 14 15 1

Write a function that takes a 4-by-4 matrix as an argument and then determine if the matrix is magic or not. The #matrix should be stored as a two-dimensional list. Test your function with a magic matrix and a none magic matrix.

def magic(matrix = []):

    magic_matrix = False
    if len(matrix) != 4:
        print('Enter a 4 * 4 matrix')
        return magic_matrix

    row1Sum = sum(matrix[0]) 
    rowSum_ok = True
    for row in range(1, 4):
        if sum(matrix[row]) != row1Sum:
            rowSum_ok = False
            break
    colSum_ok = True
    for col in range(4):
        s_col = 0
        for row in range(4):
            s_col += matrix[row][col]
        if s_col != row1Sum:
            colSum_ok = False
            break
        if rowSum_ok and colSum_ok:
            magic_matrix = True
        return magic_matrix

def mainMagic():
    m1 = [[9, 6, 3, 16],
          [4, 15, 10, 5],
          [14, 1, 8, 11],
          [7, 12, 13, 2]]
    print('\nThe matrix:')
    for i in range(4):
        for j in m1[i]:
            print(str(j).rjust(3), end =' ')
        print()
    if magic(m1):
        print('is a magic matrix.')
    else:
        print('is not a magic matrix.')

    m2 = [[6, 22, 44, 18],
          [1, 11, 10, 13],
          [3, 17, 6, 12],
          [9, 14, 2, 1]]
    print('\nThe matrix:')
    for i in range(4):
        for j in m2[i]:
            print(repr(j).rjust(3), end = ' ')
        print()
    if magic(m2):
        print('is a magic matrix.')
    else:
        print('is not a magic matrix.')

mainMagic()

Answer 1

With a couple of set comprehensions and a zip() that is fairly straight forward to cleanup like:

Code:

def is_magic(matrix):
    sum_rows = {sum(row) for row in matrix}
    sum_cols = {sum(col) for col in zip(*matrix)}
    return len(sum_cols) == 1 and sum_cols == sum_rows

Test Code:

m1 = [[9, 6, 3, 16],
      [4, 15, 10, 5],
      [14, 1, 8, 11],
      [7, 12, 13, 2]]

m2 = [[6, 22, 44, 18],
      [1, 11, 10, 13],
      [3, 17, 6, 12],
      [9, 14, 2, 1]]

print(is_magic(m1))
print(is_magic(m2))

Results:

True
False