如何将N×M正整数矩阵A转换为N×M×L矩阵B（无循环实现）

给定元素均为正整数的N×M矩阵A，需转换为N×M×L的矩阵B（其中L = 1 + A的最大值），要求当k = A[i,j]时B[i,j,k] = 1，其余位置为0。以下是几种无需循环的numpy实现方案：

方法1：利用单位矩阵索引 #

这是最简洁的实现方式，借助np.eye生成单位矩阵后直接索引：

import numpy as np

max_val = np.amax(A)
L = 1 + max_val
B = np.eye(L)[A]

原理 #

np.eye(L)生成一个L×L的单位矩阵，每一行对应一个独热向量。用矩阵A作为索引时，numpy会自动将A中每个元素映射到单位矩阵的对应行，最终直接得到形状为(N, M, L)的目标矩阵B。

方法2：利用np.put_along赋值 #

如果需要先初始化全零矩阵再赋值，可以用np.put_along完成向量化赋值：

import numpy as np

max_val = np.amax(A)
L = 1 + max_val
B = np.zeros((A.shape[0], A.shape[1], L), dtype=np.int32)
# 给A扩展一个维度，匹配B的第三维索引
indices = A[..., np.newaxis]
np.put_along(B, indices, 1, axis=2)

原理 #

先创建符合目标形状的全零矩阵，将A扩展为(N, M, 1)的形状后，作为第三维的索引数组，np.put_along会沿着指定轴（这里是第三轴）把对应位置的值设为1。

方法3：利用广播与相等判断 #

通过广播机制让A和所有可能的k值比较，生成布尔矩阵后转成整数：

import numpy as np

max_val = np.amax(A)
L = 1 + max_val
k_vals = np.arange(L)
# 扩展A的维度后和k_vals广播比较，再转成整数类型
B = (A[..., np.newaxis] == k_vals).astype(np.int32)

原理 #

将A从(N, M)扩展为(N, M, 1)，和形状为(L,)的k_vals进行广播比较，得到的布尔矩阵中，符合k = A[i,j]的位置为True，转成整数后就是1，其余为0。

以上三种方法均为numpy的向量化操作，相比循环实现，在处理大规模矩阵时效率会显著提升。

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