Numpy：使用索引数组为三维数组赋值

问题描述 #

我有一个形状为(2, 2, 3)的indices数组，具体如下：

import numpy as np
indices = np.array([[[ 0,  6, 12],
        [ 0,  6, 12]],
       [[ 1,  7, 13],
        [ 1,  7, 13]]])

需要将这些值作为索引，把np.zeros生成的矩阵对应位置设为1。已创建形状为(2, 2, 18)的one_hot零数组：

one_hot = np.zeros((2, 2, 18))

期望得到的结果如下：

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

希望用Numpy高级索引类似one_hot[indices] = 1的方式实现，该如何操作？

解决方案 #

要实现这个需求，关键是构造和indices维度匹配的前两个维度的索引数组，配合indices完成三维高级索引：

方法一：利用np.indices生成网格索引 #

# 生成前两个维度的索引数组，形状为(2,2)
i, j = np.indices(indices.shape[:2])

# 将i、j扩展为(2,2,1)，和indices的(2,2,3)维度对齐，完成赋值
one_hot[i[..., np.newaxis], j[..., np.newaxis], indices] = 1

方法二：利用np.meshgrid生成网格索引 #

# 生成前两个维度的网格索引，indexing='ij'确保维度顺序和原数组一致
i, j = np.meshgrid(np.arange(2), np.arange(2), indexing='ij')

# 同样扩展维度后赋值
one_hot[i[..., np.newaxis], j[..., np.newaxis], indices] = 1

原理说明 #

one_hot是三维数组(2,2,18)，定位元素需要三个索引：第一维索引、第二维索引、第三维索引（即indices中的值）。通过将前两个维度的索引扩展为和indices相同的三维结构，Numpy会自动广播这些索引，精准定位到每个需要置1的位置，完成赋值。

内容的提问来源于stack exchange，提问作者Max S.