How to pass list of indices along a dimension of a 2D array in numpy?

Question

Suppose I have a 2D numpy array of shape, say, 4x3:

myarray = np.array([[0,1,2],
                    [3,4,5],
                    [6,7,8],
                    [9,10,11]])

and I have a list of indices corresponding to the second dimension, with length 4 and values ranging from 0 to 2, i.e., for each of the 4 rows, I have one different index corresponding to the value I want to select from that row:

idx = [0,2,1,2]

How can I pass this list of indices to the 2D array and get as a result the following 1D array of length 4, where each element corresponds to the indexed value from each row of the original array?

array([ 0,  5,  7, 11])

I am looking for a solution that doesn't require looping as I intend to do this for very large arrays.

Answer 1

You should use zip to iterate over two arrays simultaneously.

data = [
    [1,2,3,4,5],
    [2,4,6,8,10],
    [3,6,9,12,15]
]

indexes = [0,1,2]

for (arr, i) in zip(data, indexes):
    print(arr[i])
    
# Or more pythonic way ?    
print([arr[i] for (arr,i) in zip(data, indexes)]) # [1, 4, 9]

Answer 2

For simplicity I will reduce the dimensions to 2x2 to make the example easier to show.

Suppose we have a 2D array:

arr = np.array(
  [[1,2], 
   [3,4]]
)
nrows = arr.shape[0] # 1000 in your case

and a 1D array of indexes:

idx = np.array([1,0])

In your case the 2D array will have dims 1000x40 and the 1D array of indexes dim 1000.

1. Convert indexes into a 2D array of shape 1000x1

idx = np.array([0,1]).reshape(-1,1) 

2. Use the following to select element at index in each row according to the 1D vector idx

arr[np.arange(arr.shape[0])[:,None], idx]

The np.arange(arr.shape[0])[:,None] simply generates a selector for all rows in your 2D array.

Your output will look like this.

array([[2],
       [3]])

Hope this is helpful!