基于与外部数组的交集的numpy recarray索引

时间:2011-03-10 19:01:01

标签: indexing 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,'Bob',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]

但是如果a.id有重复项,我会覆盖id_x_index中的dict值,就像在这种情况下我得到的那样

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

我知道我忽略了一些简单的方法来获得适当的指数进入重组。谢谢你的帮助。

2 个答案:

答案 0 :(得分:1)

Numpy中最简洁的方法是

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

答案 1 :(得分:1)

对于那些拥有旧版numpy的人,你也可以这样做:

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