如何在NumPy中找到另一数组元素对应的索引？

已知两个NumPy数组：

• a1：元素唯一且已排序，包含a2的所有元素
• a2：所有元素均存在于a1中

需要生成数组res，使得a1[res] == a2。

最优实现方法 #

利用np.searchsorted方法，基于二分查找快速定位元素索引，这是针对该场景最高效的方案：

import numpy as np

# 输入数组
a1 = np.array([3,4,5,7,8,9])
a2 = np.array([5,3,5,8])

# 生成结果
res = np.searchsorted(a1, a2)

# 验证
print(res)  # 输出: [2 0 2 4]
print(a1[res] == a2)  # 输出: [ True  True  True  True]

方法说明 #

np.searchsorted针对已排序数组采用二分查找逻辑，时间复杂度为O(M*logN)（N是a1长度，M是a2长度）。因为a1包含a2的所有元素，所以返回的插入位置恰好就是元素在a1中的索引，完全满足需求。

对比其他方案 #

如果用字典映射的方式实现：

idx_map = {val: idx for idx, val in enumerate(a1)}
res = np.array([idx_map[val] for val in a2])

该方法虽然可行，但对于大型数组，searchsorted的性能更优——它是NumPy内置的C级实现，避免了Python层面的循环和字典操作开销，在数据量较大时差距明显。

内容的提问来源于stack exchange，提问作者f. c.