如何遍历numpy数组的指定维度？三维数组子数组函数应用问询

你提到的把(2,2,2)三维数组视作“以一维数组为元素的二维矩阵”，本质上就是要对每个沿第3轴（索引为2）的一维子数组应用函数，下面给你几种实用的实现方式：

1. 用np.apply_along_axis（通用自定义函数场景） #

这是最直接的方法，专门用来沿指定轴对数组的子切片应用函数。先创建测试数组：

import numpy as np
# 创建(2,2,2)的三维数组
arr = np.arange(8).reshape(2,2,2)
# 数组结构：
# [[[0 1], [2 3]],
#  [[4 5], [6 7]]]

假设你有一个自定义函数，比如计算每个一维数组的乘积加1：

def my_custom_func(x):
    # x是传入的一维数组，比如[0,1]、[2,3]等
    return x[0] * x[1] + 1

然后用apply_along_axis指定轴为2（也就是每个一维子数组所在的轴）：

result = np.apply_along_axis(my_custom_func, axis=2, arr=arr)
# 输出结果是(2,2)的二维数组：
# [[1, 7],
#  [21, 43]]

这个方法的好处是不管你的自定义逻辑多复杂，只要能接收一维数组并返回单个值（或固定形状的结果），都能直接用。

2. 用numpy内置的轴操作（内置函数场景） #

如果你的函数是numpy已经实现的操作（比如求和、均值、最大值等），直接用数组的轴方法会比apply_along_axis高效得多，因为这是完全向量化的操作：

# 计算每个一维数组的均值
mean_result = arr.mean(axis=2)
# 输出：
# [[0.5, 2.5],
#  [4.5, 6.5]]

# 计算每个一维数组的求和
sum_result = arr.sum(axis=2)
# 输出：
# [[1, 5],
#  [9, 13]]

3. 用np.vectorize（带签名的向量化） #

如果你想把自定义函数包装成向量化函数，也可以用np.vectorize，但要注意指定signature参数，告诉numpy输入是一维数组，输出是标量：

vec_func = np.vectorize(my_custom_func, signature='(n)->()')
result = vec_func(arr)
# 得到和apply_along_axis一样的结果

不过要注意，np.vectorize本质上是循环的封装，性能不一定比apply_along_axis好，更适合需要重复调用这个向量化函数的场景。

总结一下：你要的“沿两个轴应用函数”其实不需要同时操作两个轴，只需要定位到一维子数组所在的轴（这里是轴2），沿这个轴应用函数就会自动得到二维的结果矩阵。

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