CODEWITHSUNDEEP
pythonnumpy
Matt
Matt

Reputation: 1839

numpy indexing operations for 3D matrix

Is there an elegant/quick way to reproduce this without the for loops? I'm looking to have a 3D matrix of values, and and 2D matrix that gives the indices for which to copy the 3rd dimensions' values while creating a new 3D matrix of the same shape. Here is an implementation with a lot of loops. 

np.random.seed(0)

x = np.random.randint(5, size=(2, 3, 4))
y = np.random.randint(x.shape[1], size=(3, 4))

z = np.zeros((2, 3, 4))

for i in range(x.shape[0]):
    for j in range(x.shape[1]):
        z[i, j, :] = x[i, y[i, j], :]

Upvotes: 0

Views: 110

Answers (1)

hpaulj
hpaulj

Reputation: 231385

This puzzled me for a bit, until I realized you aren't using all of y. y is (3,4), but you are indexing over (2,3):

In [28]: x[np.arange(2)[:,None], y[:2,:3],:]
Out[28]: 
array([[[4, 0, 0, 4],
        [4, 0, 3, 3],
        [3, 1, 3, 2]],

       [[3, 0, 3, 0],
        [2, 1, 0, 1],
        [1, 0, 1, 4]]])

We could use all of y with:

In [32]: x[np.arange(2)[:,None,None],y,np.arange(4)]
Out[32]: 
array([[[4, 0, 3, 2],
        [4, 0, 3, 2],
        [3, 0, 0, 3]],

       [[3, 1, 1, 4],
        [3, 1, 1, 4],
        [1, 1, 3, 1]]])

the 3 indexes broadcast to (2,3,4). But the selection is different from your z.

Upvotes: 1

Related Questions