基于NumPy实现二维数组指定索引赋值1的高效Pythonic方案

问题说明 #

输入NumPy数组：

import numpy as np
indexes = np.array([[4], [3], [2], [1]])

期望输出：

output = [[0, 0, 0, 0, 1], [0, 0, 0, 1, 0], [0, 0, 1, 0, 0], [0, 1, 0, 0, 0]]

需求是将输入数组中的每个元素作为索引，生成长度为输入最大值+1的数组（因索引从0开始），对应索引位置赋值为1，替代低效的for循环实现。

方法1：直接向量化赋值（最通用） #

利用NumPy的索引机制直接批量赋值，完全避免循环开销：

import numpy as np

indexes = np.array([[4], [3], [2], [1]])
max_len = indexes.max() + 1
output = np.zeros((indexes.shape[0], max_len), dtype=int)
# 生成行索引和列索引，批量赋值
output[np.arange(indexes.shape[0]), indexes.flatten()] = 1

np.arange(indexes.shape[0])生成每行的索引，indexes.flatten()将输入的二维数组转为一维列索引，通过数组索引直接完成批量赋值，性能远优于Python循环。

方法2：借助单位矩阵快速生成 #

如果输入索引是连续正整数，可直接利用单位矩阵切片：

import numpy as np

indexes = np.array([[4], [3], [2], [1]])
max_len = indexes.max() + 1
# 生成单位矩阵后，用索引值取对应行
output = np.eye(max_len, dtype=int)[indexes.flatten()]

np.eye(max_len)会生成一个max_len×max_len的单位矩阵，每一行正好对应一个独热数组，直接用输入索引取出对应行即可得到目标结果。

方法3：保留输入维度的赋值方式 #

如果不需要展平输入数组，可用np.put_along沿指定轴完成赋值：

import numpy as np

indexes = np.array([[4], [3], [2], [1]])
max_len = indexes.max() + 1
output = np.zeros((indexes.shape[0], max_len), dtype=int)
# 沿axis=1方向，将indexes指定的位置设为1
np.put_along(output, indexes, 1, axis=1)

这种方法无需修改输入数组的维度，适合需要保持原数组结构的场景。

以上三种方法均基于NumPy的向量化操作，底层由C实现，处理大型数组时效率远高于Python层面的for循环。

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