如何高效优雅实现Numpy一维数组的数值-索引映射？

这确实是处理百万级规模数组时很实用的需求，尤其是需要快速定位某个值的所有出现位置时。咱得兼顾运行效率和代码优雅性，下面给你几种经过验证的靠谱方案：

方案一：基于Numpy排序与分割（最推荐，百万级数据首选） #

这个方法完全依赖Numpy的向量化操作，避开了Python循环，性能拉满，特别适合百万级别的大数组：

import numpy as np

def map_value_to_indices(a):
    # 先对数组的索引按值排序
    sorted_indices = np.argsort(a)
    sorted_values = a[sorted_indices]
    
    # 找到值发生变化的位置，作为分割点
    split_positions = np.where(sorted_values[1:] != sorted_values[:-1])[0] + 1
    
    # 把排序后的索引按值分割成组
    index_groups = np.split(sorted_indices, split_positions)
    
    # 获取所有唯一值，对应每个索引组
    unique_values = sorted_values[np.concatenate([[0], split_positions])]
    
    # 根据需求生成输出结构：如果值是连续非负整数，生成列表；否则用字典
    max_val = unique_values.max()
    if np.array_equal(unique_values, np.arange(max_val + 1)):
        output = [[]] * (max_val + 1)
        for val, indices in zip(unique_values, index_groups):
            output[val] = indices
    else:
        output = {val: indices for val, indices in zip(unique_values, index_groups)}
    
    return output

为什么这个方法高效？ #

Numpy的argsort是底层优化过的排序实现，时间复杂度是O(n log n)，后续的分割和查找操作都是向量化的，几乎没有Python层面的开销，处理百万级数据非常快。

方案二：用字典收集索引（代码更直观） #

如果你的数组值分布比较零散，或者更看重代码的可读性，可以用collections.defaultdict来收集每个值对应的索引：

from collections import defaultdict
import numpy as np

def map_value_to_indices_dict(a):
    index_map = defaultdict(list)
    # 用np.ndenumerate遍历数组，比Python原生enumerate更适配Numpy数组
    for idx, val in np.ndenumerate(a):
        index_map[val].append(idx[0])
    
    # 可选：把列表转成Numpy数组，方便后续数值操作
    for val in index_map:
        index_map[val] = np.array(index_map[val])
    
    return index_map

优缺点 #

代码逻辑非常直观，但因为涉及Python层面的遍历，百万级数据下的速度会比方案一慢一些，但如果唯一值的数量不多，这个差距可以忽略。

方案三：用np.unique + np.where（代码最简洁） #

如果你的数据量不算特别大，或者唯一值的数量很少，这个方案的代码最简洁：

import numpy as np

def map_value_to_indices_simple(a):
    unique_values = np.unique(a)
    output = {}
    for val in unique_values:
        output[val] = np.where(a == val)[0]
    return output

注意事项 #

这个方案的问题在于：每个np.where都会遍历整个数组，时间复杂度是O(n*k)（k是唯一值数量）。当k很大时（比如数组里大部分值都是唯一的），百万级数据下的效率会明显下降，所以更适合小数据量或唯一值少的场景。

总结选择建议 #

• 若处理百万级大数组，且值是连续非负整数 → 方案一最优，性能拉满。
• 若值分布零散，看重代码可读性 → 方案二更合适。
• 若数据量小或唯一值极少 → 方案三代码最简洁。

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