如何将一维Numpy数组转换为适配二维数组索引更新的掩码数组？

给定二维Numpy数组a和一维数组b，我们需要把b中每个元素对应的行索引位置转换为掩码数组——每行仅在b指定的列位置为1，其余为0。以下是两种高效实现方法：

测试数据 #

import numpy as np

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

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

方法一：索引直接赋值 #

先创建与a形状完全一致的全0数组，再通过行索引和b提供的列索引，将指定位置设为1：

# 生成全0掩码数组
mask = np.zeros_like(a)
# 生成行索引序列（0到5，对应a的6行）
row_indices = np.arange(len(b))
# 为每行指定列位置赋值1
mask[row_indices, b] = 1

print(mask)

输出结果：

array([[0, 1, 0, 0],
       [0, 0, 1, 0],
       [0, 0, 1, 0],
       [1, 0, 0, 0],
       [0, 0, 0, 1],
       [0, 0, 0, 1]])

方法二：利用单位矩阵快速生成 #

借助np.eye()生成单位矩阵，再用b作为索引提取对应行，直接得到目标掩码：

# 根据a的列数生成单位矩阵，再用b索引取出对应行
mask = np.eye(a.shape[1], dtype=a.dtype)[b]

这种方法代码更简洁，原理是单位矩阵的第n行只有第n列是1，正好匹配我们需要的掩码格式。

得到掩码后，即可通过a[mask.astype(bool)] = 1完成原数组的更新操作。

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