我有一个矩阵,有很多行和 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. ]])
谢谢,
最佳答案
这是一种使用 np.argpartition
的矢量化方法-
m,n = a.shape
a[np.arange(m)[:,None],np.argpartition(a,n-2,axis=1)[:,:-2]] = 0
sample 运行-
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. ]])
关于python - numpy 矩阵,通过对每一行进行排序将 0 设置为值,我们在Stack Overflow上找到一个类似的问题: https://stackoverflow.com/questions/35416050/