使用NumPy查找数组匹配索引的代码错误排查及适配需求

代码错误点 #

1. 重复元素只取首次索引：你用np.where(name==array)[0][0]时，np.where会返回该元素在数组中所有匹配的索引（比如第二个'KDI'对应的结果是[0,1]），但你用[0]强行取了第一个索引，所以重复元素会输出相同的首次位置，而不是自身的索引。
2. 多余的判断逻辑：if name in array完全没必要，因为name是直接从array里遍历出来的，必然存在于数组中。

解决方案 #

分两种场景给出适配代码：

场景1：单数组获取自身索引 #

如果只是要输出每个元素对应的自身索引，直接用enumerate遍历即可：

import numpy as np
array = np.array(['KDI', 'KDI', 'KDU', 'KDA', 'ANU', 'AMU', 'BDU', 'CDU', 'CDU', 'DAI', 'DAH'], dtype='<U4')
for idx, _ in enumerate(array):
    print(idx)

运行后会输出0,1,2,...,10，完全符合你的预期。

场景2：适配双数组的匹配查找 #

如果需要处理两个不同数组的匹配（比如找数组B中每个元素在数组A中的对应索引，支持重复元素按出现顺序匹配），可以用计数器跟踪每个元素的匹配次数：

import numpy as np
from collections import defaultdict

def get_match_indices(target_arr, lookup_arr):
    # 记录每个元素已匹配的次数
    match_count = defaultdict(int)
    result = []
    for item in lookup_arr:
        # 获取当前元素在目标数组中的所有匹配索引
        all_matches = np.where(item == target_arr)[0]
        # 取对应次数的索引（第一次匹配取第0个，第二次取第1个，以此类推）
        result.append(all_matches[match_count[item]])
        match_count[item] += 1
    return result

# 测试单数组场景
array = np.array(['KDI', 'KDI', 'KDU', 'KDA', 'ANU', 'AMU', 'BDU', 'CDU', 'CDU', 'DAI', 'DAH'], dtype='<U4')
print(get_match_indices(array, array))  # 输出 [0,1,2,3,4,5,6,7,8,9,10]

# 测试双数组场景
target = np.array(['KDI', 'KDI', 'KDU', 'ANU'], dtype='<U4')
lookup = np.array(['ANU', 'KDI', 'KDI'], dtype='<U4')
print(get_match_indices(target, lookup))  # 输出 [3,0,1]

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