Python中如何提取子集向量元素在原向量中的索引？

纯Python实现（无显式for循环） #

利用字典建立元素到索引的映射，结合map()函数快速提取子集元素的索引，全程无显式for循环语句，符合Pythonic风格：

original = [3, 1, 4, 2, 5]
subset = [4, 1, 5]

# 构建元素与对应索引的映射字典
element_to_idx = dict(zip(original, range(len(original))))
# 提取子集元素的索引
indices = list(map(element_to_idx.get, subset))

print(indices)  # 输出: [2, 1, 4]

NumPy矢量化实现（适合大规模数据） #

如果处理的是大规模数组，用NumPy的矢量化操作效率更高，同样无需显式循环：

import numpy as np

original = np.array([3, 1, 4, 2, 5])
subset = np.array([4, 1, 5])

# 对原数组排序并记录排序后的原索引
sorted_indices = np.argsort(original)
sorted_original = original[sorted_indices]
# 查找子集元素在排序后数组中的位置
pos = np.searchsorted(sorted_original, subset)
# 还原为原数组的索引
indices = sorted_indices[pos]

print(indices)  # 输出: [2 1 4]

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