numpy recarray indexing based on intersection with external array

问题

I'm trying to subset the records in a numpy.recarray based on the common values between one of the recarrays fields and an external array. For example,

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])

I want to take the intersection of a.id and b to determine what records to pull from the recarray, so that I get back:

(10, 'Bob', 145.7)
(10, 'Jim', 130.5)

Naively, I tried:

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

but of course that doesn't work because there is no a[10]. I also tried to do this by creating a reverse dict between the id field and the index and subsetted from there, e.g.

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]

but then I'm overriding dict values in id_x_index if a.id has duplicates, as in this case I get

(10, 'Jim', 130.5)
(10, 'Jim', 130.5)

I know I'm overlooking some simple way to get the appropriate indices into the recarray. Thanks for help.

回答1:

The most concise way to do this in Numpy is

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

回答2:

And for those who have an older version of numpy, you can also do it this way:

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

来源：https://stackoverflow.com/questions/5264448/numpy-recarray-indexing-based-on-intersection-with-external-array