使用Numpy将二维数组每行最大值替换为1的无循环实现方案

核心方案（基于np.where实现） #

直接利用Numpy广播机制实现，全程无显式for循环，性能远高于遍历实现：
首先通过a.max(axis=1, keepdims=True)获取每行的最大值，keepdims=True参数保证输出的最大值数组维度和原数组匹配，可直接广播做逐元素对比，再通过np.where做条件替换：匹配到每行最大值的位置替换为1，其余位置保留原值。

完整示例代码 #

import numpy as np

# 示例输入数组
a = np.array([[0.5, 0.2, 0.1], 
              [0.6, 0.3, 0.8], 
              [0.3, 0.4, 0.2]])

# 核心替换逻辑
new_a = np.where(a == a.max(axis=1, keepdims=True), 1, a)

输出验证 #

>>> print(new_a)
[[1.  0.2 0.1]
 [0.6 0.3 1. ]
 [0.3 1.  0.2]]

补充说明 #

• 如果单一行内存在多个相等的最大值，上述代码会将所有最大值位置都替换为1，符合绝大多数业务场景需求。
• 如果需要仅替换每行第一个出现的最大值，可额外结合np.argmax和索引赋值实现，代码如下：

# 仅替换每行第一个最大值的实现
row_idx = np.arange(a.shape[0])
col_idx = a.argmax(axis=1)
new_a = a.copy()
new_a[row_idx, col_idx] = 1

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