在Numpy数组中查找多个值

Question

我正在寻找一个numpy函数来查找在向量(xs)中找到某些值的索引.值以另一个数组(ys)给出.返回的索引必须遵循ys的顺序.

在代码中,我想用numpy函数替换下面的列表推导.

>> import numpy as np
>> xs = np.asarray([45, 67, 32, 52, 94, 64, 21])
>> ys = np.asarray([67, 94])
>> ndx = np.asarray([np.nonzero(xs == y)[0][0] for y in ys]) # <---- This line
>> print(ndx)
[1 4]

有快速的方法吗？

谢谢

Answer 1

对于大阵列xs和ys,你将需要改变的基本方法这成为快.如果你对排序很好xs,那么一个简单的选择是使用numpy.searchsorted():

xs.sort()
ndx = numpy.searchsorted(xs, ys)

如果保持原始顺序很重要xs,您也可以使用这种方法,但是您需要记住原始索引:

orig_indices = xs.argsort()
ndx = orig_indices[numpy.searchsorted(xs[orig_indices], ys)]