更新 numpy 数组中每一行的最后一个元素

Question

我有两个 numpy arrays， array_one是 NxM 和array_two是 NxMx3，我想根据array_one的值更改array_two每一行中最后一个元素的值，如下所示：

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

其中foo基于对来自array_one的元素的计算返回一个值。

有没有办法避免手动循环 arrays 并使用 numpy 函数加快此过程？

Answer 1

示例显示使用 np.vectorize 来实现您的想法。

用你的 foo 替换 square ，你应该做生意了。

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