Index column values in 2D array using array of indices [duplicate]

Question

I have the following array:

import numpy as np 

print(A)
array([[ 0,  1,  4,  5,  8,  7],
       [ 5,  3,  4,  1,  8, 11],
       [ 2,  7,  5,  3,  4,  1],
       [ 2,  8,  8,  1, 10,  1],
       [ 2, 14,  8,  6,  5,  3]])

And I need to the values A corresponding to these column indices:

b = np.array([5, 0, 3, 4, 4])

Expected output:

array([ 7,  5,  3, 10,  5])

Thanks in advance.

Please share the expected output

yatu
– yatu

2020-04-17 13:21:29 +00:00
Commented Apr 17, 2020 at 13:21 — yatu
– yatu, Commented Apr 17, 2020 at 13:21
Edited to include :)

okvoyce
– okvoyce

2020-04-17 13:25:02 +00:00
Commented Apr 17, 2020 at 13:25 — okvoyce
– okvoyce, Commented Apr 17, 2020 at 13:25

yatu · Accepted Answer · 2020-04-17 14:01:29Z

4

You can use advanced indexing. You need to define an indexing array across the first axis, so that both indexing arrays are broadcast together and each column index refers to a specific row. In this case you just want an np.arange to index on the rows:

A[np.arange(A.shape[0]), b]
# array([ 7,  5,  3, 10,  5])

edited Apr 17, 2020 at 14:01

answered Apr 17, 2020 at 13:32

yatu

88.6k12 gold badges93 silver badges148 bronze badges

Sign up to request clarification or add additional context in comments.

Collectives™ on Stack Overflow

Index column values in 2D array using array of indices [duplicate]

1 Answer 1

Comments

Linked

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Linked

Related