Python中处理numpy数组时如何替代依赖维度索引的for循环

可以直接利用numpy的广播机制实现无循环运算，全部运算在numpy底层的C接口执行，相比Python层面的for循环性能提升非常明显，适配大体量数组的需求。

示例适配代码 #

import numpy as np
a = np.ones( (2,10) )
# 直接生成轴1的索引数组，通过广播匹配原数组形状完成运算
a *= np.arange(a.shape[1]).reshape(1, -1)

原理解释 #

• np.arange(a.shape[1]) 生成的就是你循环中用到的索引值序列[0,1,2,...,9]
• 通过reshape(1, -1)把一维索引数组调整为形状(1, 10)的二维数组，和原数组(2, 10)的形状匹配广播规则：运算时索引数组的第一维会自动复制扩展到和原数组第一维一致的长度，等价于给每一列乘对应索引值
• 最终运算结果和for循环完全一致，没有精度差异

通用扩展方法 #

如果需要依据其他维度的索引做运算，只需要调整索引数组的reshape形状即可：

• 要按第0维（行维度）的索引乘对应行所有元素：np.arange(a.shape[0]).reshape(-1, 1)
• 更高维度的数组同理，把索引数组放在对应维度的位置，其余维度填1即可适配广播规则

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