Index 3D arrays with Numpy

Question

I have an array in three dimensions (x, y and z) and an address vector. This vector has a size equal to the x dimension of the array, and its objective is for each x to point a y bringing their respective z, that is, the expected result is dimension (x, z).

Below is an example of code, which works as expected, but someone knows if you have any Numpy function with which I can replace the loop for and solve the problem more optimally?

arr = np.random.rand(100,5,2)
result = np.random.rand(100,2)
id = [np.random.randint(0, 5) for _ in range(100)]
for i in range(100): 
    result[i] = arr[i,id[i]]

Answer 1

The code below solves what you are trying to do. I compared your code with mine (setting the same Seed of the random number generator) and the results were equal.

Notice that instead of looping up to 100, I created a array 100 elements with the method np.arange(). I used this array as an index to reference the positions of arr. Also note that I created the variable id as a array, that accepts this type of advanced indexing: id[np.arange(100)].

arr = np.random.rand(100,5,2)
id = np.array([np.random.randint(0, 5) for _ in range(100)])
result = arr[np.arange(100), id[np.arange(100)]]

A tip, since you want efficiency: when creating an array and filling it later (as you did in your code), do not create using np.random.Rand(), use np.zeros(). It will save random value generation for nothing. Anyway, note that in the above code it was not even necessary, the array result was created already filled.