如何用Numpy在三列数组中通过两列匹配获取第三列值

为实现多数据库向单数据库的数据迁移，需构建内存高效的翻译表，因此用Numpy创建了包含obj_id、db_id、pk的三列结构化数组，测试数组示例如下：

import numpy as np

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

数组结构如下：

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

目标是根据(obj_id=1999994, db_id=7)映射到对应的pk值2999991，目前用低效的Python循环实现：

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

高效实现方案 #

方法一：布尔向量过滤（单次查找首选）

利用Numpy向量化操作直接生成匹配掩码，底层C实现远快于Python循环：

# 生成匹配条件的布尔掩码
mask = (a['obj_id'] == 1999994) & (a['db_id'] == 7)
# 提取第一个匹配的pk值
result = a['pk'][mask][0]
print(result)  # 输出：2999991

方法二：排序+二分查找（多次查找首选）

如果需要频繁执行查找操作，先对数组排序，之后用二分查找实现O(logn)复杂度的查询：

# 按obj_id和db_id为关键字排序数组
sorted_a = np.sort(a, order=['obj_id', 'db_id'])
# 构造匹配的键值结构
search_key = np.array((1999994, 7), dtype=[('obj_id', np.int32), ('db_id', np.int8)])
# 用二分查找定位位置
pos = np.searchsorted(sorted_a, search_key, side='left')
# 验证匹配有效性（避免键不存在的情况）
if pos < len(sorted_a) and sorted_a[pos]['obj_id'] == 1999994 and sorted_a[pos]['db_id'] == 7:
    result = sorted_a[pos]['pk']
    print(result)  # 输出：2999991
else:
    print("未找到匹配的记录")

内容的提问来源于stack exchange，提问作者Eric