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.

CodePudding user response：

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