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

我正在尝试找到一种有效的方法，让我可以将稀疏矩阵的前 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 

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

请您参考如下方法:

通过直接修改 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