Numpy:对于一个数组中的每个元素,在另一个数组中查找索引

时间:2011-11-24 02:22:05

标签: python arrays search numpy indexing

我有两个1D阵列,x& y,比另一个小。我试图在x中找到y的每个元素的索引。

我发现了两种天真的方法,第一种是缓慢的,第二种是内存密集型。

缓慢的方式

indices= []
for iy in y:
    indices += np.where(x==iy)[0][0]

记忆猪

xe = np.outer([1,]*len(x), y)
ye = np.outer(x, [1,]*len(y))
junk, indices = np.where(np.equal(xe, ye))

是否存在更快或更少内存密集型方法?理想情况下,搜索将利用这样一个事实,即我们在列表中搜索的不是一件事,而是很多事情,因此稍微更适合并行化。 如果你不假设y的每个元素实际上都是x,那么奖励积分。

8 个答案:

答案 0 :(得分:23)

我想建议一行解决方案:

indices = np.where(np.in1d(x, y))[0]

结果是一个带有x数组索引的数组,它对应于x中找到的元素。

如果需要,可以在没有numpy的情况下使用它。

答案 1 :(得分:20)

正如Joe Kington所说,searchsorted()可以非常快速地搜索元素。要处理不在x中的元素,可以使用原始y检查搜索结果,并创建一个屏蔽数组:

import numpy as np
x = np.array([3,5,7,1,9,8,6,6])
y = np.array([2,1,5,10,100,6])

index = np.argsort(x)
sorted_x = x[index]
sorted_index = np.searchsorted(sorted_x, y)

yindex = np.take(index, sorted_index, mode="clip")
mask = x[yindex] != y

result = np.ma.array(yindex, mask=mask)
print result

结果是:

[-- 3 1 -- -- 6]

答案 2 :(得分:16)

这个怎么样?

它确实假设y的每个元素都在x中,(并且即使对于没有的元素也会返回结果!)但是它要快得多。

import numpy as np

# Generate some example data...
x = np.arange(1000)
np.random.shuffle(x)
y = np.arange(100)

# Actually preform the operation...
xsorted = np.argsort(x)
ypos = np.searchsorted(x[xsorted], y)
indices = xsorted[ypos]

答案 3 :(得分:3)

numpy_indexed包(免责声明:我是它的作者)包含一个完全符合这个功能的函数:

import numpy_indexed as npi
indices = npi.indices(x, y, missing='mask')

如果并非x中的所有元素都存在于x中,它将引发KeyError;但也许我应该添加一个kwarg,以便人们可以选择用-1或其他东西标记这些项目。

它应该具有与当前接受的答案相同的效率,因为实现方式类似。然而,numpy_indexed更灵活,并且还允许搜索多维数组的行索引。

编辑:我改变了对缺失值的处理;现在可以使用'加注','忽略'或'掩码'来设置'缺失的'kwarg。在后一种情况下,您将获得一个y相同长度的掩码数组,您可以在其上调用.compressed()来获取有效索引。请注意,如果这是您需要知道的,还有npi.contains(x,y)。

答案 4 :(得分:3)

我会这样做:

indices = np.where(y[:, None] == x[None, :])[1]

与memory-hog方式不同,它使用broadcast直接生成2D布尔数组,而不为x和y创建2D数组。

答案 5 :(得分:2)

我认为这是一个更清晰的版本:

np.where(y.reshape(y.size, 1) == x)[1]

胜过indices = np.where(y[:, None] == x[None, :])[1]。您无需将x广播为2D。

我发现这种类型的解决方案是最好的,因为与在此处或其他地方发布的基于searchsorted()或in1d()的解决方案不同,以上解决方案适用于重复项,因此无需理会是否进行了排序。这对我很重要,因为我希望x遵循特定的自定义顺序。

答案 6 :(得分:0)

更直接的解决方案,并不期望对数组进行排序。

[2, 3, 3, 1, 4]

结果是:

{{1}}

答案 7 :(得分:0)

使用以下代码行:-

indices = np.where(y [:, None] == x [None,:])[1]