numpy 矩阵，通过对每一行进行排序来将值设置为 0 [英] numpy matrix, setting 0 to values by sorting each row

问题描述

我有一个矩阵，有很多行和 8 列.每个单元格代表当前行属于 8 个类别中的一个的概率.我只想保留每行中的 2 个最高值，并将其余值设置为 0.

I have a matrix, with many rows, and 8 columns. Each cell represents a probability for the current row to belong to 1 of the 8 classes. I would like to keep only the 2 highest values in each row, and set the rest to 0.

到目前为止，我能想到的唯一方法是分别对每一行进行循环和排序.例如:

So far, the only way I can think of is by looping and sorting each row separately. For example:

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.  ]])

谢谢，

推荐答案

这是一种带有 np.argpartition -

Here's one vectorized approach with 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.        ]])

这篇关于numpy 矩阵，通过对每一行进行排序来将值设置为 0的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！