如何用Numpy函数批量更新NxMx3数组每行最后一个元素？

完全可以通过Numpy的矢量化特性避免手动循环，根据你的foo函数类型分两种情况处理：

情况1：foo已支持矢量化输入 #

如果foo本身就能接收Numpy数组并返回同形状的数组（比如内部用Numpy内置函数实现），直接通过切片赋值即可：

# 利用省略号`...`匹配所有前置维度，简洁通用
array_two[..., -1] = foo(array_one)

省略号...会自动匹配array_one的所有维度（这里是N×M），精准对应array_two的前两个维度，直接替换最后一个通道的所有元素。

情况2：foo仅支持标量输入 #

如果foo是只能处理单个标量的函数，可以用np.vectorize将其包装为矢量化函数（注意：np.vectorize本质是对循环的封装，但代码更简洁，且能利用Numpy的内部优化）：

# 包装标量函数为矢量化函数
vectorized_foo = np.vectorize(foo)
# 完成赋值
array_two[..., -1] = vectorized_foo(array_one)

提示：如果追求极致性能，建议直接修改foo使其支持矢量化输入（比如用Numpy的运算替代原生Python标量运算），np.vectorize的性能提升相对有限。

示例代码 #

import numpy as np

# 示例标量函数（也可替换为矢量化版本）
def foo(x):
    return x * 2 + 1

# 创建测试数组
N, M = 3, 4
array_one = np.arange(N*M).reshape(N, M)
array_two = np.zeros((N, M, 3))

# 执行赋值
array_two[..., -1] = foo(array_one)

# 验证结果
print(array_two[1, 2, -1])  # 输出 1*4+2=6，foo(6)=13，结果应为13

内容的提问来源于stack exchange，提问作者Mdp11