2

我正在尝试根据 recarrays 字段之一和外部数组之间的公共值对 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]

但是这当然行不通,因为没有 a[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, '吉姆', 130.5)
(10, '吉姆', 130.5)

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

4

2 回答 2

1

在 Numpy 中执行此操作的最简洁方法是

subset = a[np.in1d(a.id, b)]
于 2011-03-10T19:08:27.770 回答
1

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

subset = a[np.array([i in b for i in a.id])]
于 2011-03-10T23:23:45.180 回答