Use indices of a 3d array to fill a 4d array

Question

I have two 3D arrays, one containing the values I am using and one containing indices. I want to fill a 4D array using these two.

Each entry of the index array points towards a row of the input array.4

At first I simply iterated through the values of i, j, and k and manually filled in each row. However, since this is a machine learning project, this method takes way too long.

# x.shape = (8, 2500, 3)
# ind.shape = (8, 2500, 9)

M = np.empty(8, 2500, 9, 3)

for i in range(0, M.shape[0]):
    for j in range(0, M.shape[1]):
        for k in range(0, M.shape[2]):
            M[i, j, k, :] = x[i, ind[i, j, k], :]

Is there a faster way that exists to do this?

Answer 1

You can try something like:

import numpy as np
M = x[np.arange(0,ind.shape[0])[:, None, None], ind]

where [:, None, None] is needed to broadcast np.arange(0,ind.shape[0]) to the correct dimensions for indexing the array x .

As a test, you can generate the array M with your current method, then use the above method to generate an array M_ , and confirm that (M == M_).all() returns True .

I make it to be at least 30x as fast.