CODEWITHSUNDEEP
arrayspython-3.xnumpymultidimensional-array
Allan
Allan

Reputation: 43

Indexing 3D arrays with Numpy

I have an array in three dimensions (x, y, z) and an indexing vector. This vector has a size equal to the dimension x of the array. Its objective is to index a specific y bringing their respective z, i.e., the expected result has dimension (x, z).

I wrote a code that works as expected, but does anyone know if a Numpy function can replace the for loop 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]]

Upvotes: 2

Views: 249

Answers (1)

maousi
maousi

Reputation: 46

You can achieve this with this piece of code:

import numpy as np
arr = np.random.randn(100, 5, 2)
ids = np.random.randint(0, 5, size=100)
res = arr[range(100), ids]
res.shape # (100, 2)

Upvotes: 1

Related Questions