2011-11-24 306 views
24

我有两个一维数组,其中一个比另一个小。我试图在x中找到y的每个元素的索引。Numpy:对于一个数组中的每个元素,找到另一个数组中的索引

我发现了两种天真的方法来做到这一点,第一种是慢速的,第二种是内存密集型的。

较慢的方式

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中,则为奖励分。

回答

15

正如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] 
+0

啊,这是一段很棒的代码。感谢您推广Joe的代码! – Chris

12

这个怎么样?

它确实假定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] 
+0

神奇。确实快得多。我包含了'assert na.all(na.intersect1d(x,y)== na.sort(y))'来限制输入,使得y是x的一个子集。谢谢! – Chris

15

我想建议单line solution:

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

结果是一个数组,其索引为x数组,它对应于在x中找到的来自y的元素。

可以使用它没有numpy.where如果需要。

+0

这应该是选择的答案。即使x的值重复或不存在,它也可以工作。涉及搜索的答案很复杂,很奇怪,不自然。 –

+5

虽然这确实会返回y中存在于x中的元素的索引,但返回的索引的顺序与x中的值的顺序不匹配。考虑:x = np.array([1,2,3,4,5]; y = np.array([5,4,3,2,1])。上述方法返回array([0,1, 2,3,4]),所以x [0] = 1与y [0] = 5匹配,这不是想要的... – ccbunney

+0

伟大而优雅的代码! – Orangeblue

0

numpy_indexed包(免责声明:我是它的作者)中包含的正是这样做的一个功能:

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

它将目前引发KeyError如果y中不是所有的元素存在于X;但也许我应该添加一个kwarg,以便可以选择用-1来标记这些物品。

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

编辑:ive改变了缺失值的处理; '失踪'kwarg现在可以设置'raise','ignore'或'mask'。在后一种情况下,你会得到一个相同长度的y的掩码数组,你可以在其中调用.compressed()来获得有效的索引。请注意,如果这是你所需要知道的,那么也有npi.contains(x,y)。

4

我只是这样做:

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

与你的记忆,猪的方式,这使得利用广播直接生成2D布尔数组没有创造x和y二维数组。

+1

为了记录,这使得内存成为 – romeric

+0

是的,我的意思是减少内存占用,我认为我的版本在保持代码清洁的同时占用较少的内存是一个很好的折衷方案。 –

0

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

import pandas as pd 
A = pd.Series(['amsterdam', 'delhi', 'chromepet', 'tokyo', 'others']) 
B = pd.Series(['chromepet', 'tokyo', 'tokyo', 'delhi', 'others']) 

# Find index position of B's items in A 
B.map(lambda x: np.where(A==x)[0][0]).tolist() 

结果是:

[2, 3, 3, 1, 4] 
相关问题