如何创建一个沿特定（第三）轴的 3D（shape=m,n,o）数组，索引由 2D 数组（shape=m,n）给出？答案

【问题标题】：How to create a 3D (shape=m,n,o) array with ones along a specific (third) axis, the indices are given by a 2D array (shape=m,n)?如何创建一个沿特定（第三）轴的 3D（shape=m,n,o）数组，索引由 2D 数组（shape=m,n）给出？
【发布时间】：2019-01-24 07:08:25
【问题描述】：

假设我有一个形状为 (2, 2) 的二维数组，其中包含索引

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

我想做的是创建一个形状为 (2, 2, 4) 的 3D 数组，沿第三轴的值为 1，它们的位置由x 给出，因此：

y = np.zeros(shape=(2,2,4))
myfunc(array=y, indices=x, axis=2)

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

到目前为止，我还没有找到任何索引方法。 for 循环可能可以做到这一点，但我确信有更快的矢量化方法。

【问题讨论】：

标签： python arrays numpy indexing

【解决方案1】：

您要查找的内容称为advanced indexing。要正确使用整数数组进行索引，您需要有一组广播到正确形状的数组。由于x 已经与其中两个维度对齐，因此您只需要使用沿每个轴的索引制作二维数组。 np.ogrid 对此有所帮助，因为它创建了广播到正确形状的最小范围数组：

a, b = np.ogrid[:2, :2]
y[a, b, x] = 1

ogrid的结果等价于

a = np.arange(2).reshape(-1, 1)
b = np.arange(2).reshape(1, -1)

或者

a = np.arange(2)[:, None]
b = np.arange(2)[None, :]

你也可以写一个单行：

y[(*tuple(slice(None, n) for n in x.shape), x)] = 1

【讨论】：

谁不喜欢那些晦涩难懂的蟒蛇单线！ :)
@朱利安。他们写起来很有趣，如果可能的话，我总是尝试在我的答案中发布一个。我不建议实际使用它们。
很好，这比前面提到的具有更大阵列的网格网格更快吗？
@Commander。大概有点吧。唯一的区别是它分配 2x1、1x2 而不是一对 2x2 索引。

【解决方案2】：

使用for 循环：

y = np.zeros(shape=(2,2,4))
for i in range(2):
    for j in range(2):
        y[i,j,x[i,j]] = 1

使用高级索引：

y = np.zeros(shape=(2,2,4))
i, j = np.meshgrid(np.arange(2), np.arange(2))
y[i,j,x[i,j]] = 1

输出（在这两种情况下）：

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

至于哪个会更快，你需要根据你的具体情况进行尝试......

【讨论】：

谢谢！由于我的数组会比较大，所以我倾向于使用 meshgrid 方法。