numpy matrix, установка 0 в значения путем сортировки каждой строки

Question

У меня есть матрица со многими строками и 8 столбцами. Каждая ячейка представляет вероятность того, что текущая строка принадлежит одному из 8 классов. Я хотел бы сохранить только 2 наивысших значения в каждой строке, а остальные - 0.

До сих пор, единственный способ, о котором я могу думать, - это цикл и сортировка каждой строки отдельно. Например:

a = np.array([[ 0.2 , 0.1 , 0.02 , 0.01 , 0.031, 0.11 ], 
       [ 0.5 , 0.1 , 0.02 , 0.01 , 0.031, 0.11 ], 
       [ 0.2 , 0.1 , 0.22 , 0.15 , 0.031, 0.11 ]]) 

Я хотел бы получить:

array([[ 0.2 , 0. , 0. , 0. , 0. , 0.11], 
     [ 0.5 , 0. , 0. , 0. , 0. , 0.11], 
     [ 0.2 , 0. , 0.22, 0. , 0. , 0. ]]) 

Спасибо,

Answer 1

Вот один Векторизованный подход с np.argpartition -

m,n = a.shape 
a[np.arange(m)[:,None],np.argpartition(a,n-2,axis=1)[:,:-2]] = 0 

Пример запуск -

In [570]: a 
Out[570]: 
array([[ 0.94791114, 0.48438182, 0.54574317, 0.45481231, 0.94013836], 
     [ 0.03861196, 0.99047316, 0.7897759 , 0.38863967, 0.93659426], 
     [ 0.49436676, 0.93762758, 0.33694977, 0.45701655, 0.73078113], 
     [ 0.21240062, 0.85141765, 0.00815352, 0.52517721, 0.49752736]]) 

In [571]: m,n = a.shape 
    ...: a[np.arange(m)[:,None],np.argpartition(a,n-2,axis=1)[:,:-2]] = 0 
    ...: 

In [572]: a 
Out[572]: 
array([[ 0.94791114, 0.  , 0.  , 0.  , 0.94013836], 
     [ 0.  , 0.99047316, 0.  , 0.  , 0.93659426], 
     [ 0.  , 0.93762758, 0.  , 0.  , 0.73078113], 
     [ 0.  , 0.85141765, 0.  , 0.52517721, 0.  ]]) 

Answer 2

Это должно работать, однако, она изменяет a. Это то, что вы хотите? Необходимо ли избегать циклов?

sorted = np.sort(a, axis=1) 

for idx, row in enumerate(a): 
    row[row < sorted[idx,-2]] = 0  

Или вы могли бы сделать это:

a[a < sorted[:,None,-2]] = 0