Pandas: получение имени столбца для отсортированного значения на уровне записи

Question

У меня есть датафрейм, как показано ниже:

Теперь мне нужно получить 2 верхних имени столбца максимального значения для каждой записи. Я могу сделать сортировку используя np.sort, но как мне получить имена столбцов для каждой записи? Например, для записи 0 мне нужно получить имена столбцов ['beta', 'gamma'], так как максимальное значение от gamma (4.1), а второе max от beta (2.1). И так для всех записей

Answer 1

использование Series.nlargest с np.dot

new_df = ( df.apply(lambda x: x.nlargest(2),axis=1).notna()
             .dot(df.add_prefix('_').columns)
             .str.split('_',expand=True)
             .iloc[:,1:]

          )
print(new_df)
       1      2
0   beta  gamma
1   beta  gamma
2  alpha  gamma

Answer 2

Если важна производительность, используйте numpy.argsort для позиций отсортированных значений и получения имен столбцов путем индексации:

arr = df.columns.values[np.argsort(-(df).values)[:, :2]]
print (arr)
[['gama' 'beta']
 ['beta' 'gama']
 ['gama' 'beta']]

df1 = pd.DataFrame(arr, index=df.index).rename(columns=lambda x: f'top_{x+1}')
print (df1)
  top_1 top_2
0  gama  beta
1  beta  gama
2  gama  beta

Если производительность не важна:

df1 = pd.DataFrame([x.nlargest(2).index for _, x in df.T.items()])
print (df1)
      0     1
0  gama  beta
1  beta  gama
2  gama  beta