Python中不使用Pandas，基于Numpy获取主数组中子集元素的索引

针对你需要找出subset数组每个元素在main数组中对应索引的需求，这里提供两种纯Numpy的实现方法：

方法一：高效排序映射法（适合大规模数据） #

这种方法利用排序和搜索机制，效率更高，尤其数组规模较大时表现更优：

import numpy as np

# 示例数据
main = np.array([0.3, 0.2, 0.7, 0.9])
subset = np.array([0.3, 0.9])

# 获取main数组排序后的原索引
sorted_main_indices = np.argsort(main)
# 生成排序后的main数组
sorted_main = main[sorted_main_indices]

# 在排序后的数组中定位subset元素，再映射回原索引
subset_indices = sorted_main_indices[np.searchsorted(sorted_main, subset)]

print(subset_indices)  # 输出: [0 3]

方法二：直观向量化查找法（适合小规模数据） #

这种方法逻辑更直观，代码简洁，对于你的数组规模（600和230）完全够用：

import numpy as np

# 示例数据
main = np.array([0.3, 0.2, 0.7, 0.9])
subset = np.array([0.3, 0.9])

# 定义向量化的索引查找函数
find_index = np.vectorize(lambda val: np.where(main == val)[0][0])
# 批量获取subset元素的索引
subset_indices = find_index(subset)

print(subset_indices)  # 输出: [0 3]

错误原因说明 #

你遇到的IndexError大概率是直接用subset数组作为索引去访问main数组（比如写了main[subset]），而Numpy数组的索引只能是整数、切片、布尔数组等类型，浮点数数组无法作为有效索引，因此触发错误。

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