在多个轴上应用numpy.searchsorted

Question

我正在尝试传递一个2d数组的排序值和一个值的数组,searchsorted并让它返回一个相应的1d索引值数组.我必须多次执行此操作,并且我正在尝试找到一种有效的方法.最后,我想传递一个沿着轴排序的3d数组值,一个2d数组值来搜索并让函数返回一个2d索引数组.

这里有一个 MLE

from numpy import *
random.seed(1234)

a = zeros((4,10))
a[:,1:-1] = random.rand(4,8)
a[:,1:-1].sort(1)
a[:,-1] = 1.

v = random.rand(4)

res = array([searchsorted(a[j], v[j]) for j in xrange(4)])

res应该在哪里[9,1,7,6]

有没有一种有效的方法来做到这一点？理想情况下,我想尽可能避免使用Cython.如果它有帮助,搜索的每个数组应该有9-10个元素,而搜索的值的数量更大(100-1000个元素),因此传递的2d数组的大小为1000x10.

编辑修正了MLE中的拼写错误

Answer 1

您可以更换

res = array([searchsorted(a[j], v[j]) for j in xrange(4)])

和：

res = (a.T < v).sum(axis=0)

就是这样（我会在评论中写出正确性的证明）。在上面的行中，虽然我们没有进行二分搜索，但由于我们在较大维度上对迭代进行向量化，因此我们获得了40 倍的加速（在我的 Mac 上）。这种加速是通过二维数组大小获得的(1000,10)。

为了测量执行情况，您可以使用以下命令来平均1.67100 次执行中原始命令的运行时间（毫秒）：

python2 -m timeit -n 100 \
    -s 'import numpy as np; np.random.seed(1234); a = np.zeros((1000,10)); a[:,1:-1] = np.random.rand(1000,8); a[:,1:-1].sort(1); a[:,-1] = 1.; v = np.random.rand(1000)' \
    'res = np.array([np.searchsorted(a[j], v[j]) for j in xrange(1000)])'

100 loops, best of 3: 1.67 msec per loop

而建议的执行只需要41.7微秒：

python2 -m timeit -n 100 \
    -s 'import numpy as np; np.random.seed(1234); a = np.zeros((1000,10)); a[:,1:-1] = np.random.rand(1000,8); a[:,1:-1].sort(1); a[:,-1] = 1.; v = np.random.rand(1000)' \
    'res = (a.T < v).sum(axis=0)'

100 loops, best of 3: 41.7 usec per loop