Python中高效构造元素为i²+j²的N×M矩阵的方法

架构老林

2026-8-12

用Numpy高效构造i²+j²矩阵的方案

当然可以用Numpy实现，而且效率会比Python原生循环高很多——核心是利用Numpy的向量化运算和广播机制，彻底避开两层for循环的开销。

方法一：基于广播机制的最简实现

这是效率最高的写法之一，利用Numpy自动广播不同形状数组的特性：

import numpy as np

def build_matrix(N, M):
    # 生成行索引的平方列向量（形状：N×1）
    row_sq = np.arange(N)[:, np.newaxis] ** 2
    # 生成列索引的平方行向量（形状：1×M）
    col_sq = np.arange(M) ** 2
    # 广播相加得到N×M矩阵
    return row_sq + col_sq

原理：np.arange(N)[:, np.newaxis]把一维数组转成列向量，和行向量col_sq相加时，Numpy会自动将两者扩展为N×M的矩阵后再逐元素相加，全程用C级别的底层运算，没有Python循环的额外开销。

方法二：用meshgrid生成索引矩阵

如果偏好更直观的索引生成方式，可以用np.meshgrid：

import numpy as np

def build_matrix(N, M):
    # 生成行、列索引矩阵，indexing='ij'确保i对应行、j对应列
    i, j = np.meshgrid(np.arange(N), np.arange(M), indexing='ij')
    return i ** 2 + j ** 2

这个方法生成的i和j都是N×M的矩阵，直接平方相加即可，代码可读性更强，效率和第一种方法几乎无差别。