Python 返回行中第二大值的列名

Question

数据

data = [['john', 0.20, 0.0, 0.4, 0.40],['katty', 0.0, 1.0, 0.0, 0.0],['kent', 0.0, 0.51, 0.49, 0.0]]
df = pd.DataFrame(data, columns=['name','fruit', 'vegetable', 'softdrinks', 'icecream'])
df = df.set_index('name')
df.head()

期望的结果

data = [['john', 0.20, 0.0, 0.4, 0.40,'softdrinks','icecream'],['katty', 0.0, 1.0, 0.0, 0.0,'vegetable','NaN'],['kent', 0.0, 0.51, 0.49, 0.0,'vegetable','softdrinks']]
df = pd.DataFrame(data, columns=['name','fruit', 'vegetable', 'softdrinks', 'icecream', 'max_no1', 'max_no2'])
df = df.set_index('name')
df.head()

尝试过只返回最高值列名的 idxmax，我需要找到第二高的行值列名，如何实现？

多谢

Answer 1

首先通过DataFrame.mask将0设置为缺失值，然后通过DataFrame.stack重塑，对于 top2 使用SeriesGroupBy.nlargest ，最后DataFrame.join通过DataFrame.pivot重塑数据：

df1 = df.mask(df == 0).stack().groupby(level=0, group_keys=False).nlargest(2).reset_index()
df1 = df1.assign(a = df1.groupby('name').cumcount().add(1))

df = df.join(df1.pivot('name','a','level_1').add_prefix('max_no'))
print (df)
       fruit  vegetable  softdrinks  icecream     max_no1     max_no2
name                                                                 
john     0.2       0.00        0.40       0.4  softdrinks    icecream
katty    0.0       1.00        0.00       0.0   vegetable         NaN
kent     0.0       0.51        0.49       0.0   vegetable  softdrinks

或者使用DataFrame.idxmax注释的解决方案，并通过与numpy广播进行比较再次设置缺失值：

df1 = df.mask(df == 0)
df['max_no1'] = df1.idxmax(axis=1)
m = df1.columns.to_numpy() == df['max_no1'].to_numpy()[:, None]
#pandas below 0.24
#m = df1.columns.values == df['max_no1'].values[:, None]
df1 = df1.mask(m)
df['max_no2'] = df1.idxmax(axis=1)
print (df)
       fruit  vegetable  softdrinks  icecream     max_no1     max_no2
name                                                                 
john     0.2       0.00        0.40       0.4  softdrinks    icecream
katty    0.0       1.00        0.00       0.0   vegetable         NaN
kent     0.0       0.51        0.49       0.0   vegetable  softdrinks