如何使用numpy.searchsorted按二维数组首列查找对应索引

numpy.searchsorted本身仅支持一维有序数组输入，你只需要先提取目标的首列一维序列，再传入方法即可，这是目前性能最高的向量化实现方案，不需要额外循环遍历。

实现代码 #

import numpy as np

# 示例数组
A = np.array([[1,2,3],[1,2,3],[2,2,3],[3,2,1]])
# 待查找的首列值列表
B = [1,3]

# 提取首列一维数组，直接传入searchsorted
res = np.searchsorted(A[:, 0], B)
print(res)
# 输出：[0 3]

注意事项 #

• 必须保证提取的首列A[:, 0]是升序排列，这是searchsorted的强制要求，未排序的序列会返回错误结果。
• 如果需要匹配首列中该值最后一次出现的索引，可以增加参数side='right'后再减1，示例：np.searchsorted(A[:,0], B, side='right') - 1。
• 该方案的时间复杂度为O(N + MlogN)，其中N为数组A的行数，M为待查找列表B的长度，性能远高于普通Python循环实现。

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