应用numpy.在多个轴上搜索排序

Question

我试图将一个2d数组的排序值和一个1d数组的值传递给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。

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)的二维数组的大小而增加。

为了测量执行情况，可以使用以下命令在100次执行中平均原始命令的运行时间(1.67毫秒):

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