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

Question

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

更具体地说，我想对 NxM 数组 A 的每一行应用 np.triu，然后连接所有结果。我可以通过

实现

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

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

Broadcast function that changes dimension of the input array

arrays

numpy

multidimensional-array

array-broadcasting