Accessing an array by using another array of indexes in numpy

Question

I have an array A of shape [10, 10, 10].

Now I want to access that array by using another array B of shape [10, 10, 10, 3] which contains indexes.

As the output I want to get an array C of shape [10, 10, 10]. So each index in B is "replaced" by the corresponding value in A.

Unfortunately I couldn't find an appropriate answer to solve that problem despite iterating over the index array and getting each corresponding value of A step by step for each element. I'm looking for a more efficient solution.

Thanks a lot!

Answer 1

Here are two ways to do what you want. The first uses loops, and the second doesn't. The second is faster by a factor of about 10.

soltuion 1 - loops

import numpy as np
a = np.random.normal(0,1,(10,10,10)) # original array

b = np.random.randint(0,10, (10,10,10,3)) # array of positions

c = np.zeros((10,10, 10)) # new array

for i in range(a.shape[0]):
    for j in range(a.shape[1]):
        for k in range(a.shape[2]):
            c[i,j,k]=a[tuple(b[i,j,k])]

This approach takes about 4ms on my laptop

The above can be a baseline to compare to. Now here is the same thing done with array slicing and no loops.

solution 2 - No loops. More efficient array slicing

a_original_shape = a.shape
# reshape b to make it (10**3, 3) 
# b_flat[:,0] hold all the i coords, b_flat[:,0] holds j coords etc
b_flat = b.reshape( (np.product(a_original_shape),)+(3,) )

# slice out the values we want from a. This gives us a 1D array
c2_flat = a[[b_flat[:,i] for i in range(3)]]

# reshape it back to the original form. 
# All values will be the correct place in this new array
c2 = c2_flat.reshape(a_original_shape)

This approach takes about 0.5ms on my laptop

you can check that these approaches give the same thing using

np.all(c2==c)
True

The second solution here takes about 10% of the time of the first solution