Question

让我们说我有一个二维（2，2）形状的数组，其中包含索引

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

我想做的是创建一个形状为（2，2，4）的3D数组，该数组沿第三轴的值为1，其位置由x给出，因此： / p>

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循环也许可以做到这一点，但是我确信有一个更快的矢量化方法。

Answer 1

您要寻找的被称为advanced indexing。为了正确地对整数数组建立索引，您需要具有一组广播为正确形状的数组。由于x已与两个维度对齐，因此您只需要制作2D数组，并在每个轴上都带有索引。 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

Answer 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.]]])

实际上，这会更快一些，您需要尝试一下具体情况...

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

2 个答案: