CODEWITHSUNDEEP
pythonnumpy
jackl
jackl

Reputation: 147

Get indexes from one Numpy matrix based on the values of a 1D Numpy array?

I've tried np.take but I don't think that's what I'm looking for. Here's what I want to do:

arr = np.array([[3, 4], [5, 6], [7, 8]])
indexes = np.array([0, 1, 0])

n = len(arr)
result = np.zeros(n)
for i in range(n):
    result[i] = arr[i, indexes[i]]

Sorry if I just don't understand how to use np.take and that's what it is for.

Upvotes: 1

Views: 29

Answers (1)

Ehsan
Ehsan

Reputation: 12417

Advanced indexing will also keep the dtype:

arr[np.arange(len(arr)), indexes]
#[3 6 7]

Upvotes: 1

Related Questions