我的代码可以更简单吗？

Question

如果n行的矩阵（数字表）的行与每列的总和相同，则它是一个魔术矩阵。例如，下面的4乘4矩阵是魔术矩阵。每行的总和和每列的总和完全相同，值为34。

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

编写一个以4乘4矩阵作为参数的函数，然后确定矩阵是否是魔术。该     #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

有一些set comprehensions和一个zip()可以直接清理，如：

代码：

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

测试代码：

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))

结果：

True
False