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

Question

Let us say I have a 2D array with a shape (2, 2), containing indices

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

What I would like to do is to create a 3D array with a shape (2, 2, 4), which has values 1 along the third axis, and their position is given by x , therefore:

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

So far I have not found any indexing method for this. A for loop might be able to do this, but I am sure there is a faster, vectorized method.

Answer 1

What you are looking for is called advanced indexing . To index with integer arrays properly, you need to have a set of arrays that broadcast to the right shape. Since x is already lined up with two of the dimensions, you only need to make 2D arrays with the indices along each axis. np.ogrid helps with that since it creates minimal range arrays that broadcast to the correct shape:

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

The results of ogrid are equivalent to

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

Or

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

You can also write a one-liner:

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

Answer 2

With for loops:

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

With advanced indexing:

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

Output (in both cases):

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

As for which will be actually faster, you need to try on your exact case...