如何使用Numpy高效查找数组中各非-1元素的出现索引？

给定数组及对应索引：

array = [-1, -1, -1, -1, -1, -1, 3, 3, -1, 3, -1, -1,  2,  2, -1, -1,  1, -1]
# indexes  0   1   2   3   4   5  6  7   8  9  10  11  12  13  14  15  16  17

需求 #

找出数组中每个非-1元素对应的所有出现索引，返回按元素值升序排序的索引列表（例如元素1对应[16]，元素2对应[12,13]，元素3对应[6,7,9]）。

约束条件 #

• 忽略值为-1的元素；
• 数组约含50个元素，元素值最大为500；
• 该处理逻辑需被调用6000次以上，需保证性能。

Numpy实现方案 #

利用Numpy向量化操作特性，避免Python循环，确保高效执行：

import numpy as np

def group_non_neg1_indices(arr):
    # 转换为numpy数组
    np_arr = np.array(arr)
    # 筛选非-1元素的索引与对应值
    mask = np_arr != -1
    values = np_arr[mask]
    indices = np.arange(len(np_arr))[mask]
    
    # 获取升序排列的唯一元素
    unique_vals = np.unique(values)
    
    # 按元素值分组索引并转为列表
    result = {}
    for val in unique_vals:
        result[val] = indices[values == val].tolist()
    
    return result

执行效果 #

调用函数处理示例数组：

array = [-1, -1, -1, -1, -1, -1, 3, 3, -1, 3, -1, -1,  2,  2, -1, -1,  1, -1]
print(group_non_neg1_indices(array))

输出结果：

{1: [16], 2: [12, 13], 3: [6, 7, 9]}

性能说明 #

• 全程使用Numpy向量化操作，比纯Python循环的单轮调用开销低一个数量级；
• np.unique直接返回升序唯一值，省去额外排序步骤，简化逻辑同时提升效率；
• 针对50元素规模+6000次调用的场景，该方案的执行效率完全满足要求。

内容的提问来源于stack exchange，提问作者Andrei M.