是否可以将这个numpy数组比较矢量化?

mdfafbf1  于 2021-08-20  发布在  Java
关注(0)|答案(1)|浏览(494)

我在python中有两个numpy数组:

a = np.array(sorted(np.random.rand(6)*6)) # It is sorted.
b = np.array(np.random.rand(3)*6)

假设数组是

a = array([0.27148588, 0.42828064, 2.48130785, 4.01811243, 4.79403723, 5.46398145])
b = array([0.06231266, 1.64276013, 5.22786201])

我想生成一个包含索引的数组,其中 a<= 比中的每个元素都多 b ,即我想要的正是:

np.argmin(np.array([a<b_i for b_i in b]),1)-1

产生 array([-1, 1, 4]) 意思是 b[0]<a[0] , a[1]<b[1]<a[2]a[4]<b[2]<a[5] .
是否有任何本地numpy快速矢量化的方法来避免 for

dwthyt8l

dwthyt8l1#

来回答您的具体问题,即,获取 np.array([a<b_i for b_i in b] ,您可以利用广播,在这里,您可以使用:

a[None, ...] < b[..., None]

因此:

>>> a[None, ...] < b[..., None]
array([[False, False, False, False, False, False],
       [ True,  True, False, False, False, False],
       [ True,  True,  True,  True,  True, False]])

重要的是,对于广播:

>>> a[None, ...].shape,  b[..., None].shape
((1, 6), (3, 1))

这里是官方numpy文档的链接,以了解广播。一些相关的花絮:
在两个阵列上操作时,numpy会按元素比较它们的形状。它从尾随(即最右侧)维度开始,然后向左运行。当需要时,两个维度是兼容的
他们是平等的,或者
其中一个是1
...
当比较的任一尺寸为一个时,使用另一个。换句话说,尺寸为1的尺寸标注会拉伸或“复制”以匹配其他尺寸标注。

编辑

正如您问题下的评论所指出的,使用完全不同的方法在算法上比您自己的暴力解决方案要好得多,即利用二进制搜索,使用 np.searchsorted

相关问题