如何使用NumPy内置方法将数组最后一维最大值替换为1最小值替换为0

实现方案 #

你可以直接使用NumPy向量化操作实现需求，全程不需要任何循环，代码如下：

import numpy as np

# 方法1：分步赋值，可读性更高
res = np.zeros_like(a)
# 生成判断掩码：最后一维第一个元素 >= 第二个元素时为True，正好匹配相等时第一个取1的规则
mask = a[..., 0] >= a[..., 1]
res[..., 0] = mask.astype(np.float64)
res[..., 1] = 1 - res[..., 0]

如果想要更简洁的写法，可以直接用np.where实现：

# 方法2：np.where 一行实现
res = np.where(a[..., [0]] >= a[..., [1]], [1.0, 0.0], [0.0, 1.0])

效果验证 #

用你提供的示例数组运行上述代码，输出结果和预期完全一致：

array([[[[1., 0.],
         [1., 0.],
         [1., 0.]],

        [[0., 1.],
         [1., 0.],
         [0., 1.]]]])

优势说明 #

两种实现都是完全向量化的操作，底层由C语言执行，不管数组前几个维度有多大，执行效率都远高于Python层面的嵌套循环，且写法更简洁，适配任意前导维度、最后一维长度为2的数组，不需要修改代码就能适配不同shape的输入。

内容的提问来源于stack exchange，提问作者A. Arash Chitgar