Python函数返回值规范问题及幻方检测代码优化咨询

问题说明 #

需要修正幻方检测代码的返回值问题，满足两个要求：

1. 函数有非None返回值时，必须添加显式return语句；
2. 函数中所有return语句要么都返回表达式，要么都不返回。

原代码：

import collections

import numpy as np


def is_square_magic(rows_and_cols):
    """is square magic"""

    matrix = [[int(num) for num in input().split()] for _ in range(rows_and_cols)]

    matrix_length = np.array(matrix).size
    uniqueness_check = len(np.unique(matrix))
    ascending_sort = np.unique(matrix)

    if matrix_length == uniqueness_check and ascending_sort[-1] == pow(rows_and_cols, 2):
        horizontal_flip = np.trace(matrix)
        vertical_flip = sum(np.fliplr(matrix).diagonal())
        sum_of_rows = np.sum(matrix, axis=1)
        sum_of_columns = np.sum(matrix, axis=0)

        rows_check = collections.Counter(sum_of_rows)
        columns_check = collections.Counter(sum_of_columns)

        return (horizontal_flip == vertical_flip).any() and (rows_check == columns_check)


if is_square_magic(int(input())):
    print("YES")
else:
    print("NO")

修正后的代码 #

import numpy as np


def is_square_magic(matrix):
    """判断输入矩阵是否为幻方"""
    n = len(matrix)
    # 检查矩阵是否为标准n×n结构
    if any(len(row) != n for row in matrix):
        return False
    
    flat_matrix = np.array(matrix).flatten()
    # 检查元素唯一性和数值范围是否符合幻方要求
    if len(np.unique(flat_matrix)) != n*n or flat_matrix.max() != n*n:
        return False
    
    # 计算两条对角线的和
    main_diag_sum = np.trace(matrix)
    anti_diag_sum = np.trace(np.fliplr(matrix))
    
    # 对角线和不相等则不是幻方
    if main_diag_sum != anti_diag_sum:
        return False
    
    # 计算所有行和与列和
    row_sums = np.sum(matrix, axis=1)
    col_sums = np.sum(matrix, axis=0)
    
    # 检查所有行和、列和是否与对角线和一致
    if not (np.all(row_sums == main_diag_sum) and np.all(col_sums == main_diag_sum)):
        return False
    
    # 所有条件满足，返回True
    return True


# 读取输入并执行判断
n = int(input())
input_matrix = [[int(num) for num in input().split()] for _ in range(n)]
print("YES" if is_square_magic(input_matrix) else "NO")

核心修改点 #

• 统一返回规则：函数所有分支都返回布尔值（True/False），彻底解决原代码仅部分分支有返回值的问题。
• 分离输入与逻辑：将输入读取逻辑移到函数外部，让函数只负责幻方判断逻辑，提升代码复用性和可测试性。
• 优化判断逻辑：
  • 新增矩阵结构合法性检查，避免非n×n矩阵传入；
  • 用np.all()直接判断行和、列和是否统一，比原代码的Counter方式更直观准确；
  • 移除多余的.any()调用（对角线和为标量，比较结果本身就是布尔值）。
• 清晰化注释：添加关键步骤的注释，提升代码可读性。

内容的提问来源于stack exchange，提问作者cyberMatter.