Лучшие 3 значения на строку в Pandas

Question

У меня есть большой фрейм данных Pandas, который находится в вене:

| ID | Var1 | Var2 | Var3 | Var4 | Var5 |
|----|------|------|------|------|------|
| 1  | 1    | 2    | 3    | 4    | 5    |
| 2  | 10   | 9    | 8    | 7    | 6    |
| 3  | 25   | 37   | 41   | 24   | 21   |
| 4  | 102  | 11   | 72   | 56   | 151  |
...

, и я хотел бы сгенерировать вывод, который выглядит следующим образом, принимая имена столбцов 3 самых высоких значений для каждой строки:

| ID | 1st Max | 2nd Max | 3rd Max |
|----|---------|---------|---------|
| 1  | Var5    | Var4    | Var3    |
| 2  | Var1    | Var2    | Var3    |
| 3  | Var3    | Var2    | Var1    |
| 4  | Var5    | Var1    | Var3    |
...

Я пытался использовать df.idmax (axis = 1), который возвращает имя первого максимального столбца, но не уверен, как вычислить два других?

Любая помощь по этому вопросу будет очень признателен, спасибо!

Answer 1

Используйте numpy.argsort для позиций отсортированных значений с помощью выбора top3 путем индексации, в последний раз передайте его конструктору DataFrame:

df = df.set_index('ID')
df = pd.DataFrame(df.columns.values[np.argsort(-df.values, axis=1)[:, :3]], 
                  index=df.index,
                  columns = ['1st Max','2nd Max','3rd Max']).reset_index()
print (df)
   ID 1st Max 2nd Max 3rd Max
0   1    Var5    Var4    Var3
1   2    Var1    Var2    Var3
2   3    Var3    Var2    Var1
3   4    Var5    Var1    Var3

Или, если производительность не важнаиспользуйте nlargest с apply для каждой строки:

c = ['1st Max','2nd Max','3rd Max']
df = (df.set_index('ID')
        .apply(lambda x: pd.Series(x.nlargest(3).index, index=c), axis=1)
        .reset_index())