Python如何获取pandas DataFrame每行最大值对应的列序号

核心实现 #

直接使用pandas内置向量化操作，全程无循环，性能最优，支持处理超大规模数据集：

import pandas as pd

# 假设你的DataFrame变量名为df
res = (
    # 第一步：排除不需要的日期列，仅保留数值列
    df.drop("Date", axis=1)
    # 第二步：转numpy数组按行取第一个最大值的索引（默认从0开始计数），加1转为1开始的序号
    .to_numpy()
    .argmax(axis=1) + 1
    # 第三步：转为列表格式得到结果
).tolist()

效果验证 #

用你给出的示例数据测试：

# 构造示例数据
data = {
    "Date": ["01.01.2020", "02.01.2020", "03.01.2020", "02.01.2020"],
    "Company 1": [23, 22, 11, 2],
    "Company 2": [21, 12, 11, 14],
    "Company 3": [14, 22, 12, 3]
}
df = pd.DataFrame(data)

# 执行代码
print(res)

输出结果为：[1, 1, 3, 2]，与期望结果完全匹配。

特殊场景适配 #

如果你的数据中存在空值NaN，请使用以下兼容版本，自动跳过空值取最大值：

res = (
    df.drop("Date", axis=1)
    # 按行取第一个最大值对应的列名，自动跳过NaN
    .idxmax(axis=1, skipna=True)
    # 把列名转为1开始的列序号，不依赖列名规则
    .apply(lambda col: df.drop("Date", axis=1).columns.get_loc(col) + 1)
).tolist()

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