Numpy：如何将 f: (形状 (3) ndarray) --> (float) 应用于形状 (...,3) 的ndarray以获得形状 (...) 的ndarray？

Question

我有一个函数，它可以把一个形状为(3)的数组转换成一个浮点数。而我还有一个形状为(...,3)的数组。请问，怎样才能把这个函数应用到那个数组上，从而得到一个形状为(...)的数组呢？

谢谢。

Answer 1

你需要用到 numpy.apply_along_axis 这个功能。

def f(a):
     return a[0] + a[1] + a[2]
mm = numpy.random.randn(5, 3)
numpy.apply_along_axis(f, 1, mm)

输出结果是：array([-1.75875289, -0.34689792, 0.66092486, -0.21626001, -0.14125476])