numpy 矩阵，通过对每一行进行排序将值设置为 0答案

【问题标题】：numpy matrix, setting 0 to values by sorting each rownumpy 矩阵，通过对每一行进行排序将值设置为 0
【发布时间】：2016-05-26 18:11:03
【问题描述】：

我有一个矩阵，有很多行和 8 列。每个单元格代表当前行属于 8 个类别中的 1 个类别的概率。我想只保留每行中的 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.  ]])

谢谢，

【问题讨论】：

标签： python sorting numpy matrix

【解决方案1】：

这应该可以，但是它会改变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

【讨论】：

认为您需要a < sorted[:,None,-2] 来保持2D 形状。因此，它只是a[a < sorted[:,None,-2]] = 0，忽略平局。
@Divakar 我不得不承认，奇怪的是我不知道使用None 进行索引，这是非常有用的知识 - 谢谢！
好吧，我上瘾了！希望你能更频繁地使用它:)
当然可以:)

【解决方案2】：

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

【讨论】：