如何将标签numpy数组无循环单行实现转换为目标二值矩阵

需求 #

将存储分类标签的一维numpy数组转换为二值矩阵，规则如下：

• 矩阵每一行对应一个唯一标签
• 原数组位置的标签和行对应标签一致时元素值为1，否则为0
• 不使用显式循环语句，优先单行实现
• 行顺序和标签在原数组中首次出现的顺序保持一致

测试输入：

import numpy as np
a = np.array([355, 355, 1005, 7005, 7005, 7005])

期望输出：

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

单行实现代码 #

完全基于numpy广播机制实现，无显式循环，代码如下：

m = (a.reshape(-1, 1) == a[np.sort(np.unique(a, return_index=True)[1])]).T.astype(int)

运行后输出结果和期望完全一致。

逻辑说明 #

• np.unique(a, return_index=True)[1]：获取每个唯一标签在原数组中第一次出现的索引
• np.sort(...)：将首次出现索引按从小到大排序，保证后续取到的唯一标签顺序和原数组中的出现顺序一致
• a[...]：按排序后的索引取值，得到按首次出现顺序排列的唯一标签数组
• a.reshape(-1, 1)：将原一维标签数组转换为列向量，形状为(原数组长度, 1)
• ==：利用numpy广播特性，列向量和唯一标签行向量做批量相等比较，直接得到形状为(原数组长度, 唯一标签数)的布尔矩阵，全程不需要手动写循环
• .T：对布尔矩阵做转置，得到形状为(唯一标签数, 原数组长度)的矩阵，此时每一行对应一个标签的匹配结果
• .astype(int)：将布尔类型的True/False转换为整数1/0，得到最终的二值矩阵

如果你确认所有标签的数值大小顺序和首次出现顺序完全一致（比如本次测试用例的场景），可以使用更简短的简化版本，但该版本会按标签数值大小排列行顺序，遇到顺序不一致的场景会输出错误结果：

m = (a[:, None] == np.unique(a)).T.astype(int)

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