百万级不同长度数组中匹配元素位置的高效查询方案

由于array_2已排序且元素唯一，同时你的数据规模达到百万级，二分查找是最优选择，相比循环调用np.where的暴力匹配，时间复杂度从O(n*m)降到O(n log m)，性能提升显著。

实现代码 #

import numpy as np

array_1 = np.array([555, 641, 1000, 641, 4, 641])
array_2 = np.array([4, 555, 641, 1000])

# 利用二分查找直接获取索引
out = np.searchsorted(array_2, array_1).tolist()
print(out)  # 输出: [1, 2, 3, 2, 0, 2]

核心原理 #

np.searchsorted专门针对有序数组实现二分查找，为array_1的每个元素计算其在array_2中保持有序的插入位置。因为array_2元素唯一且array_1的所有元素都存在于array_2中，这个插入位置就等于元素在array_2中的实际索引。

补充方案（针对未排序的array_2） #

如果遇到array_2未排序的场景，可以先构建值到索引的映射字典，再批量查询：

value_index_map = {val: idx for idx, val in enumerate(array_2)}
out = [value_index_map[x] for x in array_1]

不过对于已排序的array_2，np.searchsorted的性能比字典映射更优，尤其在数据量极大时。

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