How to create a new list where new_array[i][j] = b[a[i][j]] (with a being an array and b a vector) w-CodePudding

I have two arrays, for example a = np.array([[0, 2, 0], [0, 2, 0]]) and b = np.array([1, 1, 2]). What I want to do is to create a new array with the same size of a, but where each entry (i,j) corresponds to the value of list b with the index given by a[i][j]. Formally, I want new_list[i][j] = b[a[i][j]].

I know that this can be achieved with for loops, as shown in the code below. However, I wanted to ask if this is possible to do without for loops and only with Numpy or Python built-in functions using code vectorization.

a = np.array([[0, 2, 0], [0, 2, 0]])
b = np.array([0, 0, 2])
new_array = np.empty((2,3))
for i in range(len(a)):
    for j in range(3):
        new_array[i][j] = b[a[i][j]]

expected output:

array([[0, 2, 0],
       [0, 2, 0]])

CodePudding user response：

You can use numpy.take:

np.take(b, a)

output:

array([[0, 2, 0],
       [0, 2, 0]])

non ambiguous example

a = [[0, 2, 0], [1, 1, 2]]
b = [6, 7, 8]
np.take(b, a)

# array([[6, 8, 6],
#        [7, 7, 8]])