python - 在 csr_matrix 中保留每行 6 个最大的元素

Question

我有一个 Python scipy csr_matrix 如下。

A = csr_matrix(np.random.choice(a=[0, 1, 2, 3, 4], 
                                p=[0.35, 0.2, 0.15, 0.15, 0.15], 
                                size=[10, 12]))

密集表示将是:

[[0 2 1 2 1 0 2 0 4 2 1 2]
 [0 0 1 0 2 1 0 0 0 1 4 0]
 [1 3 3 2 1 1 3 0 0 4 2 0]
 [4 0 0 0 0 1 1 2 1 3 0 3]
 [3 0 3 1 1 3 0 3 4 4 4 0]
 [0 4 0 3 0 4 4 4 0 0 3 2]
 [0 3 0 0 3 0 1 0 3 0 0 0]
 [0 2 0 1 2 0 4 1 3 2 1 0]
 [0 2 0 4 1 1 1 3 4 2 1 1]
 [0 2 3 0 3 0 4 2 3 0 4 1]]

现在，我希望每行保留六个最大的元素，并将其余元素设置为零。如果有多个元素相等，选择哪一个留下，哪一个置零并不重要，只要每行只有六个非零元素即可。我们可以说索引较低的元素保留下来。上述矩阵的预期结果将是(手动制作):

[[0 2 1 2 0 0 2 0 4 2 0 2]
 [0 0 1 0 2 1 0 0 0 1 4 0]
 [1 3 3 2 0 0 3 0 0 4 2 0]
 [4 0 0 0 0 1 1 2 0 3 0 3]
 [3 0 3 0 0 3 0 0 4 4 4 0]
 [0 4 0 3 0 4 4 4 0 0 3 0]
 [0 3 0 0 3 0 1 0 3 0 0 0]
 [0 2 0 1 2 0 4 0 3 2 0 0]
 [0 2 0 4 1 0 0 3 4 2 0 0]
 [0 2 3 0 3 0 4 0 3 0 4 0]]

我可以想出一种通过遍历行来实现此目的的方法，但所讨论的实际矩阵很大，因此我宁愿将循环保持在最低限度。有没有人知道如何开始解决这个问题？

编辑

以下代码片段完全符合我的要求，但效率极低。以这种方式转换 3000 x 300 矩阵需要 9.3 秒。

from timeit import timeit
import warnings

from scipy.sparse import csr_matrix, SparseEfficiencyWarning

import numpy as np


__updated__ = '2015-03-12'


def random_csr_matrix():
    np.random.seed(1)
    A = csr_matrix(np.random.choice(a=[0, 1, 2, 3, 4],
                                    p=[0.35, 0.2, 0.15, 0.15, 0.15],
                                    size=[3000, 300]))
    return A

def keep_top_N_in_csr_matrix(A, N):
    warnings.simplefilter('ignore', SparseEfficiencyWarning)

#     print ('\nBefore:')
#     print (A.todense())


    N_rows = A.shape[0]
    for i in range(N_rows):
        row = np.squeeze(np.asarray(A.getrow(i).todense()))
        A[i, :] = keep_top_N_in_np_array(row, N)

#     print ('\nAfter:')
#     print (A.todense())

def keep_top_N_in_np_array(A, N):
    assert (A >= 0).all(), 'All elements shall be nonnegative'

    for _ in range(N):
        i_max = A.argmax()
        A[i_max] = -A[i_max]
    A = np.array([-i if i < 0 else 0 for i in A])

    return A

def doit_once():
    A = random_csr_matrix()
    keep_top_N_in_csr_matrix(A, 6)

if __name__ == '__main__':
    print (timeit(doit_once, number=1))

Answer 1

这是对csr_matrix 中的每一行进行操作的解决方案:

import numpy as np
from scipy.sparse import *
import scipy as sp

A = csr_matrix(
    np.random.choice(a=[0, 1, 2, 3, 4],
                     p=[0.35, 0.2, 0.15, 0.15, 0.15],
                     size=[10, 12])
)

B = A.copy()
print A.todense()

for i in range(10):
    # Get the row slice, not a copy, only the non zero elements
    row_array = A.data[A.indptr[i]: A.indptr[i+1]]
    if row_array.shape[0] <= 6:
        # Not more than six elements
        continue

    # only take the six last elements in the sorted indeces
    row_array[np.argsort(row_array)[:-6]] = 0

print 'A_mod: ', A.todense()
print 'B: ', B.todense()
print B.todense() - A.todense()

在循环的每次迭代中，您得到的不是每一行的副本，而是一个引用。因此，对 row_array 的修改也会更改稀疏矩阵 A 中的相应行。您可以调整循环中的行操作以满足更具体的标准。

更新:

在其他答案的启发下，使用 argsort 或 sort。进行此调整后，稀疏矩阵中将仅保留六个元素。