фильтрация кадра данных на основе порогового значения с известными максимальными столбцами и ограничениями порядка столбцов

Question

У меня есть фрейм данных формы (7985, 3048), отформатированный примерно так:

image_name    col1_ID    col2_ID    col3_ID    ...  colN_ID      # header
rdr_001_101   0.00342    0.00015    0.10094    ...  0.34210      # prob. distrib
rdr_001_133   0.00432    0.00025    0.01465    ...  0.22108      # prob. distrib
...
...
rdr_003_167   0.07543    0.00576    0.07523    ...  0.73421      # prob. distrib

Теперь для каждого изображения я хочу выбрать только 100 столбцов и соответствующие им индексы из этих 3047 значений. , 100 столбцов следует выбирать на основе значений в столбцах (т.е. выбирать столбцы с максимальными значениями, пока мы не достигнем значения 100).

Ограничение : порядок столбцов имеет значение (т. Е. Он должен оставаться в том же порядке на выходе, что и на входе). Пожалуйста, посмотрите пример ввода / вывода ниже для примера.

У меня есть решение на основе al oop (с использованием df.iterrows), но оно очень медленное и кажется, что это не рекомендуемый метод в целом. Кроме того, это простой пример, в то время как в моем фактическом фрейме данных может быть больше строк.

Каковы другие возможные (более быстрые) альтернативы?

---

Пример ввода, только с 5 столбцами:

image_name    col1_ID    col2_ID    col3_ID   col4_ID   col5_ID
rdr_001_101   0.00342    0.00015    0.10094   0.34210   0.27651
rdr_001_133   0.43211    0.00025    0.01465   0.22108   0.00123
rdr_003_167   0.07543    0.25761    0.07523   0.01148   0.00341

Ожидаемый результат с 3 столбцами в качестве порога:

col3_ID    col4_ID    col5_ID
col1_ID    col3_ID    col4_ID
col1_ID    col2_ID    col3_ID

Answer 1

Используйте numpy.argsort с массивом 1d по длине столбцов, сначала по np.arange(len(df.columns)), затем индексируйте верхние значения N, сортируйте их по np.sort и по последним именам столбцов cols:

df = df.set_index('image_name')
cols = df.columns.to_numpy()
N = 3

print (np.arange(len(df.columns)))
[0 1 2 3 4]

print (np.arange(len(df.columns))[np.argsort(-df.to_numpy())[:, :N]])
[[3 4 2]
 [0 3 2]
 [1 0 2]]


arr = np.sort(np.arange(len(df.columns))[np.argsort(-df.to_numpy())[:, :N]], axis=1)
print (arr)
[[2 3 4]
 [0 2 3]
 [0 1 2]]

c = cols[arr]
print (c)
[['col3_ID' 'col4_ID' 'col5_ID']
 ['col1_ID' 'col3_ID' 'col4_ID']
 ['col1_ID' 'col2_ID' 'col3_ID']]

Последнее преобразование вывода в DataFrame при необходимости:

c1 = [f'top{x+1}' for x in np.arange(N)]
df1 = pd.DataFrame(c, index=df.index, columns=c1)
print (df1)
                top1     top2     top3
image_name                            
rdr_001_101  col3_ID  col4_ID  col5_ID
rdr_001_133  col1_ID  col3_ID  col4_ID
rdr_003_167  col1_ID  col2_ID  col3_ID

---

Если возможно, сортировка имен столбцов вспомогательная 1d array не требуется:

df = df.set_index('image_name')
cols = df.columns.to_numpy()
c = np.sort(cols[np.argsort(-df.to_numpy())[:, :N]], axis=1)
print (c)

[['col3_ID' 'col4_ID' 'col5_ID']
 ['col1_ID' 'col3_ID' 'col4_ID']
 ['col1_ID' 'col2_ID' 'col3_ID']]