arrays - 改变输入数组维度的广播函数

Question

给定一些f接受1D数组并给出2D数组的函数，是否可以将它有效地应用于NxM数组的每一行A？

更具体地说，我想申请数组np.triu的每一行，然后连接所有结果。我可以通过NxMA

B = np.dstack(map(np.triu, A))

这给出了MxMxN矩阵。然而，这对于大 N 来说不是很有效。不幸的是，np.apply_along_axis这里不能使用该函数，因为f会改变维度。

了解 NumPy 在高效广播方面的强大功能后，我几乎可以肯定存在更好的解决方案来解决我的问题。

Answer 1

这是一种使用矢量化的方法broadcasting-

Bout = A.T*(np.tri(A.shape[1],dtype=bool).T[...,None])

运行时测试和输出验证 -

In [319]: A = np.random.randint(0,20,(400,100))

In [320]: %timeit np.dstack(map(np.triu, A))
10 loops, best of 3: 69.9 ms per loop

In [321]: %timeit A.T*(np.tri(A.shape[1],dtype=bool).T[...,None])
10 loops, best of 3: 24.8 ms per loop

In [322]: B = np.dstack(map(np.triu, A))

In [323]: Bout = A.T*(np.tri(A.shape[1],dtype=bool).T[...,None])

In [324]: np.allclose(B,Bout)
Out[324]: True