如何对NumPy数组最后两个维度应用函数？求替代for循环的内置方法

你完全不需要手动写for循环，NumPy提供了多种更简洁的方式来处理这类需求，分场景说明：

1. 内置线性代数函数直接批量处理 #

像np.linalg.pinv这类NumPy线性代数工具本身就支持多维数组的批量操作。当你传入(2,3,3)的数组时，函数会自动把每个a[i,:,:]这样的二维子数组单独处理，直接返回维度匹配的结果：

import numpy as np
a = np.arange(18).reshape(2,3,3)
result = np.linalg.pinv(a)

输出的result就是(2,3,3)，和你循环实现的结果完全一致，这是最高效的方式。

2. 元素级通用函数：直接用广播 #

如果你的函数是元素级操作（比如lambda x: x**3 +61），NumPy的运算符天生支持广播，直接对整个数组操作就行：

a = np.arange(18).reshape(2,3,3)
f = lambda x: x**3 + 61
result = f(a)

这会逐个元素计算，输出维度保持(2,3,3)，比循环快得多。

3. 自定义子数组操作：列表推导式或np.apply_along_axis #

如果你的函数必须作用于每个二维子数组（不是元素级），最简洁的写法是用列表推导式（语法比手动循环清爽，效率差不多）：

a = np.arange(18).reshape(2,3,3)
result = np.array([np.linalg.pinv(sub) for sub in a])

要是想更贴近"内置函数"的用法，可以用np.apply_along_axis，不过需要先把二维子数组展平处理，之后再恢复形状：

def process_subarr(arr):
    mat = arr.reshape(3,3)
    return np.linalg.pinv(mat).flatten()

a = np.arange(18).reshape(2,3,3)
result = np.apply_along_axis(process_subarr, axis=(1,2), arr=a).reshape(2,3,3)

4. 用np.vectorize包装自定义函数 #

如果需要把一个仅支持二维数组的函数改成批量版本，可以用np.vectorize，但要注意它本质是循环的语法糖，没有性能提升：

vec_process = np.vectorize(np.linalg.pinv, signature='(n,m)->(n,m)')
a = np.arange(18).reshape(2,3,3)
result = vec_process(a)

signature参数用来指定输入输出的形状，确保函数正确识别每个二维子数组。

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