2

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.

grovduck
  • 406
  • 5
  • 13

2 Answers2

1

The most concise way to do this in Numpy is

subset = a[np.in1d(a.id, b)]
Sven Marnach
  • 574,206
  • 118
  • 941
  • 841
1

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])]
joris
  • 133,120
  • 36
  • 247
  • 202
  • Thanks, that works as well. I was trying to do it with a list comprehension, but obviously wasn't quite getting the syntax correct. I'd upvote you if I could – grovduck Mar 14 '11 at 22:23