Python如何通过if语句遍历三维NumPy数组实现条件运算

• 首先代码存在基础笔误：numpy生成连续序列的函数是np.arange，不是np.arrange
• 你写的for a in arr只遍历了三维数组的第一维，每次循环拿到的a是形状为(5,5)的二维子数组，根本不是单个元素
• 直接用形状不匹配的数组做==判断会触发numpy广播规则，返回的是布尔数组，不能直接放在if语句里做单条件判断，逻辑完全不成立
• 额外注意：判断元素是否属于目标切片，本质是判断元素的索引是否落在切片范围内，不是判断元素值和切片里的值是否相等——如果数组其他位置有和切片内值相同的元素，值判断会出现误判。

方法1：显式遍历+if判断（符合使用if语句遍历的要求） #

三层循环遍历每个元素的三维索引，匹配到目标索引范围就乘0.1，其余乘1即可：

import numpy as np

arr = np.arange(125).reshape(5, 5, 5)
result = np.empty_like(arr, dtype=np.float64)

# 分别遍历三个维度的索引
for i in range(arr.shape[0]):
    for j in range(arr.shape[1]):
        for k in range(arr.shape[2]):
            # 匹配arr[0,1,:3]的索引规则：第一维=0，第二维=1，第三维取前3个位置
            if i == 0 and j == 1 and k < 3:
                result[i, j, k] = arr[i, j, k] * 0.1
            else:
                result[i, j, k] = arr[i, j, k]

方法2：numpy向量化实现（推荐，运行效率远高于手写循环） #

不需要写显式循环和if判断，利用numpy的切片赋值直接生成系数矩阵做运算即可：

import numpy as np

arr = np.arange(125).reshape(5, 5, 5)
# 先生成和arr同形状的全1系数矩阵
coef = np.ones_like(arr, dtype=np.float64)
# 把目标切片位置的系数改成0.1
coef[0, 1, :3] = 0.1
# 逐元素相乘直接得到结果
result = arr * coef

两种写法最终结果完全一致：arr[0,1,:3]位置的三个元素（值为5、6、7）会变成0.5、0.6、0.7，其余元素保持原值不变。

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