CODEWITHSUNDEEP
pythonarraysnumpyindexing
MichaelSB
MichaelSB

Reputation: 3181

Numpy: How to index 2d array with 1d array?

I have a 2d array:

a = np.random.randint(100, size=(6, 4))
[[72 76 40 11]
 [48 82  6 87]
 [53 24 25 99]
 [ 7 94 82 90]
 [28 81 10  9]
 [94 99 67 58]]

And a 1d array: 

idx = np.random.randint(4, size=6)
[0, 3, 2, 1, 0, 2]

Is it possible to index the 2d array so that the result is:

a[idx]
[72, 87, 25, 94, 28, 67]

Upvotes: 6

Views: 7591

Answers (4)

Louis Lv
Louis Lv

Reputation: 16

An alternative answer.

(a*np.eye(6,4)[idx]).sum(axis=1)

Upvotes: 0

noahn
noahn

Reputation: 31

You can also use np.diagonal if you want to avoid np.arrange.

 a = np.array([[72, 76, 40, 11],
              [48, 82,  6, 87],
              [53, 24, 25, 99],
              [ 7, 94, 82, 90],
              [28, 81, 10,  9],
              [94, 99, 67, 58]])

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

index into each array in 2d array using idx

>>> a_idx = a[...,idx]
>>> a_idx

array([[72, 11, 40, 76, 72, 40],
       [48, 87,  6, 82, 48,  6],
       [53, 99, 25, 24, 53, 25],
       [ 7, 90, 82, 94,  7, 82],
       [28,  9, 10, 81, 28, 10],
       [94, 58, 67, 99, 94, 67]])

diagonal is where the position of idx and each array in 2d array line up

>>> np.diagonal(a_idx)

array([72, 87, 25, 94, 28, 67])

Upvotes: 1

John Titor
John Titor

Reputation: 572

Is there any way to get by this without arange? It seems counterintuitive to me that something like

a[idx.reshape(-1,1)]

or

a[:,idx]

would not produce this result.

Upvotes: 3

cs95
cs95

Reputation: 402263

Since you have the column indices, all you need are the row indices. You can generate those with arange.

>>> a[np.arange(len(a)), idx]
 array([72, 87, 25, 94, 28, 67])

Upvotes: 8

Related Questions