Индексирование numy recarray на основе пересечения с внешним массивом

Question

Я пытаюсь разместить записи в numpy.recarray на основе общих значений между одним из полей recarrays и внешним массивом. Например,

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, чтобы определить, какие записи извлечь из повторного массива, чтобы я вернулся:

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

но тогда я переопределяю значения dict в id_x_index, если a.id имеет дубликаты, как в этом случае я получаю

(10, "Джим", 130,5)
(10, "Джим", 130,5)

Я знаю, что пропускаю какой-то простой способ получить соответствующие индексы в массиве. Спасибо за помощь.

Answer 1

Самый краткий способ сделать это в Numpy это

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

Answer 2

И для тех, у кого более старая версия numpy, вы также можете сделать это следующим образом:

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