如何将二维NumPy数组转换为二维独热编码？含重复值示例

给定如下二维NumPy数组：

import numpy as np
a = np.array([[0,6],
              [3,7],
              [5,5]])

需要将其转换为指定的独热编码数组，要求子数组内存在重复值时，对应行仅在该值的位置设置一个1：

b = np.array([[1,0,0,0,0,0,1,0],
              [0,0,0,1,0,0,0,1],
              [0,0,0,0,0,1,0,0]])

实现方法 #

以下是高效的NumPy实现方案：

import numpy as np

a = np.array([[0,6],
              [3,7],
              [5,5]])

# 计算独热编码的长度：基于数组中的最大值+1（元素从0开始）
one_hot_dim = a.max() + 1

# 初始化全0的结果数组
result = np.zeros((a.shape[0], one_hot_dim), dtype=int)

# 遍历每一行，去重后设置对应位置为1
for idx, row in enumerate(a):
    unique_elements = np.unique(row)
    result[idx, unique_elements] = 1

print(result)

输出结果 #

[[1 0 0 0 0 0 1 0]
 [0 0 0 1 0 0 0 1]
 [0 0 0 0 0 1 0 0]]

核心细节 #

• np.unique(row)：对每行元素去重，避免重复值多次标记1
• 索引赋值result[idx, unique_elements] = 1：直接定位到需要设1的位置，操作简洁高效
• 指定dtype=int：确保输出为整数类型，与示例格式一致

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