Zanam Zanam - 2 months ago 12
Python Question

Pulling elements from 2D array based on index array

I have an array as:

import numpy as np
A = np.arange(15).reshape(3, 5)


I also have an index array as:

ind = np.asarray([1,2,0,2,2])


The elements of ind represent the row number of A for each column of A.

i.e

I want to pull
ind[0] = 1
element from column 0 of A
I want to pull
ind[4] = 2
element from column 4 of A

Desired output is:

5, 11, 2, 13, 14

Answer

Using Numpy's fancy-indexing -

A[ind,np.arange(ind.size)]