Update last element of each row in numpy array-CodePudding

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?

CodePudding user response：

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[:,:,:]=}')