如何快速获取NumPy数组中多个指定值的全部索引？

需求：针对numpy数组a中的每个元素，找出其在numpy数组b中对应的所有索引，要求实现方式高效。

示例输入：

import numpy as np
a = np.array([1, 2, 4])
b = np.array([1, 2, 2, 3, 10, 4, 4, 4])

预期输出：[[0], [1,2], [5,6,7]]

方法一：直观遍历（适合小规模数据） #

直接遍历a中的每个元素，用np.where找出其在b中的所有索引：

result = [np.where(b == val)[0].tolist() for val in a]
print(result)

这种写法简单易懂，但如果a和b的规模较大，多次调用np.where会重复遍历b，效率较低。

方法二：批量分组映射（适合大规模数据） #

先对b进行一次分组处理，记录每个值对应的所有索引，再通过字典映射快速获取a中元素的结果，避免重复遍历b：

# 对b的元素分组，建立值到索引列表的映射
unique_vals, indices = np.unique(b, return_inverse=True)
index_groups = np.split(np.argsort(indices), np.bincount(indices).cumsum())[:-1]
val_to_indices = {val: group.tolist() for val, group in zip(unique_vals, index_groups)}

# 映射a中的元素得到结果
result = [val_to_indices[val] for val in a]
print(result)

如果a中存在b里没有的元素，可以改用val_to_indices.get(val, [])避免KeyError。

内容的提问来源于stack exchange，提问作者Valentin Macé