Python中如何通过1D NumPy数组索引生成指定位置为1的2D NumPy数组

你需要的功能属于独热编码转换，NumPy原生支持两种高效实现，均不需要手动循环，性能远高于Python层的手写逻辑，适合高频迭代和大数组场景。

两种推荐实现方式 #

方式1：np.eye索引法（代码最简洁） #

利用单位矩阵的行天然就是对应索引的独热向量的特性，直接用输入数组作为行索引提取结果：

import numpy as np

# 你的输入参数
input_arr = np.array([1, 2, 5, 1, 2])
num_classes = 6  # 预定义的2D数组宽度

# 核心转换代码
result = np.eye(num_classes, dtype=np.int32)[input_arr]

该方式单句即可完成转换，适合num_classes数值不大的场景。

方式2：高级索引赋值法（内存效率更高） #

如果预定义的宽度非常大，生成完整单位矩阵会占用额外内存，可采用先初始化零数组再批量赋值的方式：

import numpy as np

input_arr = np.array([1, 2, 5, 1, 2])
num_classes = 6

result = np.zeros((len(input_arr), num_classes), dtype=np.int32)
result[np.arange(len(input_arr)), input_arr] = 1

该方式内存占用仅和最终输出数组一致，适合num_classes较大的场景。

两种方式的输出均和你给出的示例完全匹配，单批百万级数据的转换耗时都在毫秒级。

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