Получение индексов для 3-х верхних значений из строки фрейма данных (с использованием быстрой реализации)

#python #pandas #dataframe

#python #pandas #фрейм данных

Вопрос:

Я хочу извлечь индексы для 3-х самых высоких значений для каждой строки в фрейме данных pandas. Прямо сейчас я использую

 top3df = df.apply(lambda x: pd.Series(x.nlargest(3).index), axis=1)

К сожалению, эта функция довольно дорогостоящая и выполняется на моем примере набора данных из 2 000 000 строк x 80 столбцов в течение примерно 30 минут.
Есть ли более быстрый способ?

Ответ №1:

вы можете использовать np.sort с axis=1, использовать [:,::-1] для изменения порядка сортировки, а затем [:,:3] для выбора первых 3 столбцов массива. Затем воссоздайте фрейм данных

 #input
import numpy as np

np.random.seed(3)
df = pd.DataFrame(np.random.randint(0,100,100).reshape(10, 10), 
                  columns=list('abcdefghij'))

# sort
top3 = pd.DataFrame(np.sort(df, axis=1)[:, ::-1][:,:3])
print(top3)
    0   1   2
0  74  72  56
1  96  93  81
2  90  90  69
3  97  79  62
4  94  78  64
5  85  71  63
6  99  91  80
7  96  95  61
8  91  90  74
9  88  60  56

РЕДАКТИРОВАТЬ: OP изменил вопрос, чтобы извлечь имена столбцов из 3-х верхних значений в строке, что можно сделать с argsort помощью и нарезки имен столбцов:

 print(pd.DataFrame(df.columns.to_numpy()
                     [np.argsort(df.to_numpy(), axis=1)][:, -1:-4:-1]))

1. np.sort(df, axis=1)[:,-1:-4:-1] может работать лучше, но менее читабельно.

2. Из любопытства я только что протестировал обе версии, используя фрейм данных того же размера, что и OP и %timeit магический метод ipython, и нарезка один или два раза в основном одинакова с точки зрения производительности. Хотя дважды сокращенная версия имела меньший std. dev. 18,7 мс против 49,7 мс.

3. Спасибо. К сожалению, мне также нужен x.nlargest(3).index. Поэтому у меня это не сработает. Моя ошибка в описании, я его обновил.

4. @Andreas затем вы можете использовать argsort и нарезать имена столбцов следующим образом df.columns.to_numpy()[np.argsort(df.to_numpy(), axis=1)]