Увеличение значения верхних 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