How to create an NumPy array based on the index stored in another array?

Question

Let say I have this NumPY array

A =
array([[0, 1, 3],
       [1, 2, 4]])

I have another array

B =
array([[10, 41, 26, 50, 12, 24],
       [20, 15, 42, 40, 41, 62]])

I wanted to create another array, where it selects the element in B using the index of the column in A. That is

C =
array([[10, 41, 50],
       [15, 42, 41]])

Answer 1

Try:

B[[[0],[1]], A]

Or more generally:

B[np.arange(A.shape[0])[:,None], A]

Output:

array([[10, 41, 50],
       [15, 42, 41]])

Answer 2

You can use np.take_along_axis

np.take_along_axis(B, A, axis=1)

output:

array([[10, 41, 50],
       [15, 42, 41]])

Answer 3

This can be simply done using list rather than numpy Though, in the ending we can convert it into numpy.

Code:

import numpy as np
#to make it simpler take a 1d list

a = [0,1,3]
b = [10, 41, 26, 50, 12, 24]

c = []

a = np.array(a)
b = np.array(b)


#here we are using for loop to find the value in a and append the index of b in c
for i in range(len(a)):
  print(i)
  i = a[i]
  c.append(b[i])

print(c)

c = np.array(c)
print(type(c))

#To make it more fun, you can use the random module to get random digits