如何将Pandas DataFrame每行最大值对应的列名存入新增列

实现方案 #

你可以直接使用pandas内置的idxmax()方法实现需求，该方法默认返回指定轴上最大值对应的索引值，按行取值时指定axis=1即可：

import pandas as pd
# 测试数据
test = pd.DataFrame({'0 to 4': [0.031, 0.23, 0.13], '5 to 9':[0.32, 0.142, 0.532], '10 to 14': [0.24, 0.131, 0.564]}, index=['Barking and Dagenham', 'Barnet', 'Bexley'])
# 新增列存储每行最大值对应的列名
test['max_column'] = test.idxmax(axis=1)

特殊场景适配 #

如果某行存在多个数值相同的最大值，idxmax()默认会返回第一个出现的列名。如果需要返回所有最大值对应的列名，可以用apply结合条件筛选实现：

# 多个最大值时用逗号分隔返回所有对应列名
test['max_column'] = test.apply(lambda x: ','.join(x[x == x.max()].index), axis=1)

内容的提问来源于stack exchange，提问作者Elliot.L