假设某些函数f
接受1D
数组并提供2D
数组,是否可以有效地为NxM
数组{{1}的每一行应用它}?
更具体地说,我想对A
数组np.triu
的每一行应用NxM
,然后连接所有结果。
A
给出B = np.dstack(map(np.triu, A))
矩阵。但是,对于大N来说,这不是非常有效。遗憾的是,此处无法使用函数MxMxN
,因为np.apply_along_axis
会更改维度。
了解NumPy在高效广播方面的强大功能,我几乎可以肯定,对我的问题有更好的解决方案。
答案 0 :(得分:2)
这是使用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