使用指定数字列表做for循环遍历numpy数组报IndexError如何解决

错误根因 #

你混淆了数组第一轴的索引值和待乘的系数值：

• 形状为(3,30,10)的numpy数组第一轴长度为3，合法索引只有0、1、2三个值
• 你的循环直接遍历系数列表[2,4,6]，把系数同时作为下标访问数组，第一次循环i=2时还符合索引范围，第二次i=4已经超出第一轴的最大索引2，因此抛出IndexError。

修正方案 #

方案1：调整循环逻辑 #

通过zip将索引和对应的系数一一配对遍历即可：

import numpy as np

data = np.random.rand(3,30,10)
data_new = np.zeros(data.shape)

# 索引0对应乘2、索引1对应乘4、索引2对应乘6
for idx, coef in zip(range(3), [2,4,6]):
    data_new[idx, :, :] = data[idx, :, :] * coef

方案2：使用numpy广播特性（无循环，性能更优） #

利用numpy的自动维度匹配特性，无需循环即可完成批量计算：

import numpy as np

data = np.random.rand(3,30,10)
coef = np.array([2,4,6]).reshape(-1, 1, 1)
data_new = data * coef

将(3,)形状的系数数组reshape为(3,1,1)后，会自动和(3,30,10)的输入数组维度匹配，逐元素相乘的结果和循环写法完全一致。

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