Получить имена столбцов Pandas из номеров столбцов

Question

У меня есть dataframe pandas, и я хотел бы вернуть имена столбцов с тремя самыми высокими значениями. Например:

import numpy as np 
import pandas as pd 

a = np.array([[2., 1., 0., 5., 4.], [6., 10., 7., 1., 3.]]) 
df = pd.DataFrame(a, columns=['A', 'B', 'C', 'D', 'E']) 

Дает:

A B C D E 
0 2 1 0 5 4 
1 6 10 7 1 3 

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

A B C D E First Second Third 
0 2 1 0 5 4  D  E  A 
1 6 10 7 1 3  B  C  A 

Я дошли до использования argpartition, чтобы получить индексы для трех верхних столбцов в каждой строке:

inx = df.apply(np.argpartition, args=(-3,), axis=1).ix[:, -3:].values 

Что тогда нужно, чтобы получить отсортированный

sorted_inx = inx.sort() 

Это не ясно, как я бы тогда взять эти индексы столбцов, получить имена, а затем заполнить их обратно в df как три новых столбцов

Answer 1

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

В этом случае, если вы получите индексы трех верхних значений с помощью аргумента numpy, применяемого к строкам, результирующие индексы могут быть объединены с атрибутом столбцов фрейма данных для получения результатов, которые вы ищете.

pd.concat((df, pd.DataFrame(df.columns[np.argsort(df.values, axis=1)[:, -3:][:, ::-1]], 
      columns=['First', 'Second', 'Third'])), axis=1) 

    A B C D E First Second Third 
0 2 1 0 5 4  D  E  A 
1 6 10 7 1 3  B  C  A 

Хотя повышение производительности мало для данного примера из-за накладные расходы:

>>> %timeit pd.concat((df, pd.DataFrame(df.columns[np.argsort(df.values, axis=1)[:, -3:][:, ::-1]], columns=['First', 'Second', 'Third'])), axis=1) 
100 loops, best of 3: 1.33 ms per loop 

>>> %timeit df.apply(lambda x: pd.Series(x.sort_values(ascending=False).index[:3]), axis=1) 
100 loops, best of 3: 3.55 ms per loop 

при масштабировании проблемы до улучшения становится существенным с методом применить захват 1,500x больше всего 20000 строк:

a = np.array([[2., 1., 0., 5., 4.], [6., 10., 7., 1., 3.]] * 10000) 
df = pd.DataFrame(a, columns=['A', 'B', 'C', 'D', 'E']) 

>>> %timeit pd.concat((df, pd.DataFrame(df.columns[np.argsort(df.values, axis=1)[:, -3:][:, ::-1]], columns=['First', 'Second', 'Third'])), axis=1) 
100 loops, best of 3: 4.14 ms per loop 

>>> %timeit df.apply(lambda x: pd.Series(x.sort_values(ascending=False).index[:3]), axis=1) 
1 loops, best of 3: 9.47 s per loop