Python如何高效生成取值由索引决定的Numpy数组

你写的双层嵌套循环速度慢的核心原因是所有计算都在Python解释器层面逐元素执行，完全没有用到numpy底层C实现的批量向量化运算能力。用numpy自带的坐标生成接口搭配广播机制，可以把计算速度提升3个数量级左右，完全不需要写循环。

基础实现（适配你给出的示例函数） #

直接用np.indices一次性生成目标数组所有位置对应的x、y坐标矩阵，再直接做批量运算即可，代码如下：

import numpy as np
xlength = 1920
ylength = 1080

# 生成和目标数组形状一致的坐标矩阵：第一个矩阵存行索引（对应你的y值），第二个存列索引（对应你的x值）
y_mat, x_mat = np.indices((ylength, xlength), dtype=np.float64)
# 直接向量化计算，全程在C层执行，无Python层循环开销
indexarray = x_mat ** 2 + y_mat ** 2

这段代码的计算结果和你原来的双层for循环完全一致，1920×1080尺寸的数组在普通消费级CPU上计算耗时仅需几毫秒。

通用自定义函数适配 #

只要你的自定义函数f(x,y)是基于numpy原生运算编写、支持数组形状的广播输入，都可以直接套用上述逻辑，不需要修改结构。比如你需要计算sin(x) * log(y+1)这类组合运算，直接写：

indexarray = np.sin(x_mat) * np.log(y_mat + 1)

注意：如果你的函数里有基于单个数值的条件判断逻辑，不要用Python原生的if/else，替换成np.where、np.select这类numpy提供的向量化判断接口即可，依然可以保持高速批量计算。

你之前尝试的方法为什么无效 #

• map()是Python原生的迭代方法，本质还是逐元素在Python层执行计算，和手写for循环没有速度差异，甚至会因为额外的函数调用开销更慢
• numpy.apply_over_axes、numpy.vectorize都只是封装了循环的语法糖，内部依然没有脱离Python层逐元素调用的逻辑，速度提升非常有限，远不如原生向量化实现

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