如何基于where条件与numpy数组argmax结果创建带额外维度的数组？

实现逻辑 #

• 对输入数组按指定维度计算argmax索引，标记每个维度段内的最大值位置
• 生成与原数组同形状的布尔掩码，标记每个位置是否为对应维度上的最大值
• 沿新增的最后一维对掩码、掩码取反结果做堆叠，实现最大值位置填充[1,0]、非最大值位置填充[0,1]的效果

通用实现代码 #

import numpy as np

def argmax_to_binary_map(arr: np.ndarray, arg_axis: int) -> np.ndarray:
    """
    对输入数组按指定维度计算argmax结果，映射为带额外维度的二值数组
    :param arr: 输入numpy数组
    :param arg_axis: 计算argmax的目标维度
    :return: shape比输入数组多1维（最后一维长度为2）的二值数组，最大值位置对应[1,0]，其余为[0,1]
    """
    # 计算指定维度的argmax索引
    arg_idx = np.argmax(arr, axis=arg_axis)
    # 生成与指定维度长度一致的坐标序列，用于广播比对
    axis_coord = np.arange(arr.shape[arg_axis])
    # 给argmax索引插入维度到目标arg_axis位置，广播生成同形状的最大值掩码
    max_mask = np.expand_dims(arg_idx, axis=arg_axis) == axis_coord
    # 沿最后新维度堆叠掩码、反向掩码，得到最终结果
    return np.stack([max_mask, ~max_mask], axis=-1)

示例验证 #

针对题目给出的测试输入执行调用：

# 输入数组
a = np.array([[1, 0], [1, 2], [1, 3]])
# 沿第1个维度（列方向，维度索引为1）计算映射结果，转换为整数类型0/1数值
b = argmax_to_binary_map(a, arg_axis=1).astype(np.int32)

最终输出与要求的结果完全匹配：

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

       [[0, 1],
        [1, 0]],

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

补充说明：该实现完全复用NumPy广播机制，无逐元素循环，计算效率较高；存在并列最大值时，np.argmax默认返回第一个出现的最大值位置作为标记点，若需要将所有并列最大值都标记为有效位置，可将掩码生成逻辑替换为max_mask = arr == np.max(arr, axis=arg_axis, keepdims=True)。

内容的提问来源于stack exchange，提问作者Leonardo Kanashiro Felizardo