使用np.where结合循环查找数组匹配索引的错误排查

问题背景 #

给定numpy数组：

import numpy as np
array = np.array(['KDI', 'KDI', 'KDU', 'KDA', 'ANU', 'AMU', 'BDU', 'CDU',
       'CDU', 'DAI', 'DAH', 'DBD', 'DDT', 'DDT', 'DDB', 'DDB',
       'UDM', 'UDM', 'DMA', 'DMH', 'DSX'], dtype='<U4')

执行以下代码：

for name in array:
    if name in array :
        idx_lookup = np.where(name==array)[0][0]
        print(idx_lookup)

得到结果：

0
0
2
3
4
5
6
7
7
9
10
11
12
12
14
14
16
16
18
19
20

期望输出完整索引序列[0,1,2,...,20]，且代码能适配不同数组的匹配索引查找，需定位代码错误。

错误原因 #

• 核心错误：np.where(name==array)[0][0]只会返回当前元素第一次出现的索引。比如第二个KDI执行name==array会匹配到索引0和1，但你取了结果数组的第一个元素[0]，导致重复元素都输出首次出现的索引，而非自身位置。
• 冗余判断：if name in array完全没必要，因为name是从array中遍历出来的，必然存在于数组中，这个判断对逻辑没有任何帮助。

正确实现方式 #

场景1：输出每个元素对应的自身索引 #

直接遍历数组索引即可，简单高效：

for idx in range(len(array)):
    print(idx)

或者用enumerate同时获取索引和元素（更符合Python风格）：

for idx, _ in enumerate(array):
    print(idx)

场景2：通用适配不同数组，按元素值查找所有匹配索引 #

如果需要针对任意数组，查找每个元素值对应的所有索引，可使用以下方式：

unique_names = np.unique(array)
for name in unique_names:
    all_indices = np.where(array == name)[0]
    print(f"元素 '{name}' 的所有索引: {all_indices}")

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