改变输入数组尺寸的广播功能

Question

假设某些函数f接受1D数组并提供2D数组，是否可以有效地为NxM数组{{1}的每一行应用它}？

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

我可以做到这一点

A

给出B = np.dstack(map(np.triu, A))矩阵。但是，对于大N来说，这不是非常有效。遗憾的是，此处无法使用函数MxMxN，因为np.apply_along_axis会更改维度。

了解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