增加稀疏矩阵中前k个元素的值

Question

我试图找到一种有效的方法,让我通过一些常数值增加稀疏矩阵的前k值.我目前正在使用以下代码,这对于非常大的矩阵来说非常慢:

a = csr_matrix((2,2)) #just some sample data
a[1,1] = 3.
a[0,1] = 2.

y = a.tocoo()
idx = y.data.argsort()[::-1][:1] #k is 1
for i, j in izip(y.row[idx], y.col[idx]):
    a[i,j] += 1

实际上排序似乎很快,问题在于我的最后一个循环,我通过索引排序索引来增加值.希望有人知道如何加快速度.

Answer 1

通过直接修改a.data而不是迭代行/列索引并修改单个元素,你可能会加速很多事情:

idx = a.data.argsort()[::-1][:1] #k is 1
a.data[idx] += 1

这也节省了从CSR转换 - > COO.

更新

正如@WarrenWeckesser正确地指出的那样,因为你只对k最大元素的索引感兴趣并且你不关心他们的顺序,你可以使用argpartition而不是argsort.当a.data它很大时,这可以快得多.

例如:

from scipy import sparse

# a random sparse array with 1 million non-zero elements
a = sparse.rand(10000, 10000, density=0.01, format='csr')

# find the indices of the 100 largest non-zero elements
k = 100

# using argsort:
%timeit a.data.argsort()[-k:]
# 10 loops, best of 3: 135 ms per loop

# using argpartition:
%timeit a.data.argpartition(-k)[-k:]
# 100 loops, best of 3: 13 ms per loop

# test correctness:
np.all(a.data[a.data.argsort()[-k:]] == 
       np.sort(a.data[a.data.argpartition(-k)[-k:]]))
# True