2011-03-10 39 views
2

我想基于一个recarrays字段和外部数组之间的公共值在numpy.recarray中记录子集。例如,numpy基于与外部阵列的交点的recarray索引

a = np.array([(10, 'Bob', 145.7), (20, 'Sue', 112.3), (10, 'Jim', 130.5)], 
     dtype=[('id', 'i4'), ('name', 'S10'), ('weight', 'f8')]) 
a = a.view(np.recarray) 

b = np.array([10,30]) 

我想借此a.id和B的交集,以确定哪些记录从recarray拉,让我回来:

(10,“鲍勃”,145.7 )
(10 '吉姆',130.5)

天真,我想:

common = np.intersect1d(a.id, b) 
subset = a[common] 

当然,这并不工作,因为没有一个[10]但是。我也试图通过在ID字段和索引之间创建一个反向词典并从那里进行子集化来实现这一点,

id_x_index = {} 
ids = a.id 
indexes = np.arange(a.size) 
for (id, index) in zip(ids, indexes): 
    id_x_index[id] = index 

subset_indexes = np.sort([id_x_index[x] for x in ids if x in b]) 
print a[subset_indexes] 

但后来我在id_x_index覆盖字典值如果a.id有重复的,因为在这种情况下,我得到

(10 '吉姆',130.5)
(10“吉姆',130.5)

我知道我忽略了一些简单的方法来获得适当的索引进入recarray。感谢帮助。

回答

1

最简洁的方式与NumPy做,这是

subset = a[np.in1d(a.id, b)] 
+0

优秀的,这是对尽可能简洁。谢谢 – grovduck 2011-03-10 19:29:40

1

而对于那些谁拥有numpy的一个旧版本,你也可以这样来做:

subset = a[np.array([i in b for i in a.id])] 
+0

谢谢,这也适用。我试图用列表理解的方式来做,但显然没有得到正确的语法。如果可以的话,我会让你满意的 – grovduck 2011-03-14 22:23:38