如何使用Numpy查找父数组中对应去重子数组元素的所有位置索引

首先导入numpy库并定义输入数组：

import numpy as np
A = np.array([2,1,2,1,4,4,3])
B = np.array([1,2,3,4]) # 为A去重后得到的数组

基础实现（适合中小规模数组） #

1. 输出扁平索引序列 #

result_list = np.concatenate([np.where(A == val)[0] for val in B]).tolist()
# 输出结果：[1, 3, 0, 2, 6, 4, 5]

2. 输出字典格式结果 #

result_dict = {val: np.where(A == val)[0].tolist() for val in B}
# 输出结果：{1: [1, 3], 2: [0, 2], 3: [6], 4: [4, 5]}

高性能实现（适合大规模数组，减少循环开销） #

利用排序和二分查找替代逐元素遍历，大数组下性能提升明显：

# 获得A排序后的索引与排序后数组
sort_idx = A.argsort()
sorted_A = A[sort_idx]
# 查找B中每个元素在排序数组中出现的左右边界
left_bound = np.searchsorted(sorted_A, B, side='left')
right_bound = np.searchsorted(sorted_A, B, side='right')

# 生成扁平索引序列
perf_list = np.concatenate([sort_idx[l:r] for l,r in zip(left_bound, right_bound)]).tolist()
# 生成字典格式结果
perf_dict = {B[i]: sort_idx[left_bound[i]:right_bound[i]].tolist() for i in range(len(B))}

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