0

I have a numpy array

a = np.random.random((10,3))

and I want to use b to index a

# array([0, 2, 1, 2, 0, 0, 2, 0, 0, 2])
b = np.random.choice(a=a.shape[1], size=a.shape[0])

I expect the result to be of shape (10, ), but I have tried the following and they all failed

print(a[b].shape) # (10, 3)
print(a[b,:].shape) # (10, 3)
print(a[:,b].shape) # (10, 10)
print(a[b.reshape(10,1)].shape) # (10, 1, 3)
print(a[b.reshape(1,10)].shape) # (1, 10, 3)

What would be correct way to index a in my case?

fatpanda2049
  • 483
  • 1
  • 4
  • 9

1 Answers1

1

Try using np.arange(a.shape[0]) as row index, with the array b as column index:

>>> a[np.arange(a.shape[0]), b].shape
(10,)
Sayandip Dutta
  • 15,602
  • 4
  • 23
  • 52