使用numpy如何在2列上搜索并获得3列数组中的第三个值?

rm5edbpk  于 2023-05-22  发布在  其他
关注(0)|答案(1)|浏览(94)

我正在构建一个转换表,以便将一些数据从多个数据库迁移到一个数据库。为此,我希望使用numpy的3列数组来提高内存效率。目标是将(obj_id,db_id)转换成pk。为了做一个测试,我创建了这样一个数组:

a=np.array(
    [(i*2,i % 10,i*3) for i in range(1_000_000)], 
    dtype=[('obj_id',np.int32),('db_id',np.int8),('pk',np.int32)]
)

a数组看起来像这样:

array([(      0, 0,       0), (      2, 1,       3),
       (      4, 2,       6), ..., (1999994, 7, 2999991),
       (1999996, 8, 2999994), (1999998, 9, 2999997)],
      dtype=[('obj_id', '<i4'), ('db_id', 'i1'), ('pk', '<i4')])

现在我想把(1999994, 7)翻译成2999991
在非常未优化的python中,我会这样做:

for rec in a:
   if (rec[0], rec[1]) == (1999994, 7):
     print(rec[2])
     break

我怎么能只使用numpy呢?

fv2wmkja

fv2wmkja1#

如果你有一个structured array,你可以用以下方法来切片:

a['pk'][(a['obj_id'] == 1999994) & (a['db_id'] == 7)]

输出:array([2999991], dtype=int32)

相关问题