对比三类概率列表定义argmax规则实现样本按分布归类

按指定规则的向量化实现 #

完全遵循你给出的三类分配规则，沿用二分类的布尔掩码索引思路，执行效率远高于遍历循环：

import numpy as np

# 首先将L1、L2、L3转换为numpy数组（如果原始输入是列表的话）
L1 = np.array(L1)
L2 = np.array(L2)
L3 = np.array(L3)

# 生成三类样本对应的布尔掩码
mask1 = (L1 > L2) & (L1 > L3)  # 符合规则1的样本索引
mask3 = (L1 > L2) & (L1 < L3)  # 符合规则2的样本索引
mask2 = ~(mask1 | mask3)       # 剩余符合规则3的样本索引

# 直接索引X得到分类结果
New_list1 = X[mask1]
New_list2 = X[mask2]
New_list3 = X[mask3]

通用广义argmax实现（适配任意k分类场景） #

如果规则仅为「归属到概率最大的对应分类」，可以用更通用的写法，无需手动为每个分类写掩码：

# 将所有概率列表堆叠为 shape=(样本数, 分类数) 的概率矩阵
prob_matrix = np.column_stack([L1, L2, L3])
# 取每一行最大值的索引（0对应分布1，1对应分布2，2对应分布3）
max_idx = prob_matrix.argmax(axis=1)

# 按索引分组
New_list1 = X[max_idx == 0]
New_list2 = X[max_idx == 1]
New_list3 = X[max_idx == 2]

注：无概率相等的场景下，该通用实现的输出和上述按指定规则写的实现结果完全一致。如果存在概率相等的场景，通用argmax会默认取索引更小的分类，和你指定的规则逻辑有差异，此时请优先使用第一种手动定义掩码的实现，100%匹配规则要求。

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