Python逐行比较数组扩展argmax功能 按Y_A/Y_B/Y_C最大值拆分X到A/B/C

原代码问题排查 #

• Python无for each语法，遍历索引直接写for i in range(len(X))即可
• 赋值对象错误：需求是存储X[i]到目标数组，原代码错误将整个Y_A/Y_B/Y_C数组赋值给目标变量
• 分支逻辑错误：第二个判断应使用elif，否则会和后续else形成错误的分支覆盖
• 无元素追加逻辑：空数组需要执行追加操作，不能直接覆盖赋值

正确实现 #

1. 循环实现（适配小数据量） #

import numpy as np

# 先用列表做临时存储，追加效率更高
A, B, C = [], [], []

for i in range(len(X)):
    ya, yb, yc = Y_A[i], Y_B[i], Y_C[i]
    if ya > yb and ya > yc:
        A.append(X[i])
    elif yb > ya and yb > yc:
        B.append(X[i])
    else:
        C.append(X[i])

# 转换为numpy数组和输入格式对齐
A = np.array(A)
B = np.array(B)
C = np.array(C)

2. 向量化实现（适配大数据量，性能更优） #

利用numpy布尔索引跳过显式循环，处理大规模数据时效率提升显著：

import numpy as np

# 生成分类掩码
mask_a = (Y_A > Y_B) & (Y_A > Y_C)
mask_b = (Y_B > Y_A) & (Y_B > Y_C)
mask_c = ~(mask_a | mask_b)

# 按掩码提取对应行
A = X[mask_a]
B = X[mask_b]
C = X[mask_c]

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