Update last element of each row in numpy array

Question

I have two numpy arrays, array_one which is NxM and array_two which is NxMx3, and I'd like to change the value of the last element in each row of array_two , based on values from array_one , like this:

array_two[i, j, -1] = foo(array_one[i,j])

where foo returns a value based on a computation on an element from array_one .

Is there a way to avoid manually looping over the arrays and speed up this process using numpy functions?

Answer 1

Example showing use of np.vectorize to achieve what you had in mind.

replace square with your foo and you should be in business.

import numpy as np

array_3d = np.ones((2,3,2))

array_2d = np.random.randn(2,3)

def square(x):
    return x**2

square_all = np.vectorize(square)

array_3d[:,:,-1] = square_all(array_2d)
print(f'{array_3d[:,:,:]=}')