получение индексов n max абсолютных значений в строке данных

Question

предположим, что я создаю Pandas DataFrame, как показано ниже

import pandas as pd
import numpy as np

np.random.seed(0)
x = 10*np.random.randn(5,5)
df = pd.DataFrame(x)

, например, это может сгенерировать следующее:

для каждой строки я ищу способ легко получить индексы, соответствующие наибольшим n (скажем, 3) значениям в терминах абсолютных значений.например, для первого ряда я бы ожидал [0,3,4].мы можем предположить, что результаты не должны быть упорядочены.

Я пытался найти решения, похожие на idxmax и argmax, но, похоже, они не могут легко обрабатывать несколько значений

Answer 1

Попробуйте это (это не оптимальный код):

idx_nmax = {}
n = 3
for index, row in df.iterrows():
    idx_nmax[index] = list(row.nlargest(n).index)

в конце вы получите словарь с:

в качестве ключа индекса строки

и as Values ​​индекс наибольшего значения 'n' в этой строке

Answer 2

Вы можете использовать np.argsort(axis=1)

Данный набор данных:

x = 10*np.random.randn(5,5)
df = pd.DataFrame(x)
           0          1         2          3          4
0  17.640523   4.001572  9.787380  22.408932  18.675580
1  -9.772779   9.500884 -1.513572  -1.032189   4.105985
2   1.440436  14.542735  7.610377   1.216750   4.438632
3   3.336743  14.940791 -2.051583   3.130677  -8.540957
4 -25.529898   6.536186  8.644362  -7.421650  22.697546

df.abs().values.argsort(1)[:, -3:][:, ::-1]
array([[3, 4, 0],
       [0, 1, 4],
       [1, 2, 4],
       [1, 4, 0],
       [0, 4, 2]])