查找 scipy 稀疏矩阵的一行中的前 n 个值

Question

我有一个 CSR 格式的 scipy 稀疏矩阵。它的尺寸为 72665x72665，因此将此矩阵转换为稠密矩阵来执行运算是不切实际的（此矩阵的稠密表示大约为 40 gig）。该矩阵是对称的，有大约 8200 万个非零项 (~1.5%)。

我想要做的是，对于每一行，我想要获得最大 N 值的索引。如果这是一个 numpy 数组，我会这样做np.argpartition：

    for row in matrix:
        top_n_idx = np.argpartition(row,-n)[-n:]

对于稀疏矩阵，我可以做类似的事情吗？

Answer 1

改进@Paul Panzer 的解决方案。现在它可以处理任何行的值少于 n 的情况。

def top_n_idx_sparse(matrix, n):
    """Return index of top n values in each row of a sparse matrix."""
    top_n_idx = []
    for le, ri in zip(matrix.indptr[:-1], matrix.indptr[1:]):
        n_row_pick = min(n, ri - le)
        top_n_idx.append(
            matrix.indices[
                le + np.argpartition(matrix.data[le:ri], -n_row_pick)[-n_row_pick:]
            ]
        )
    return top_n_idx

它能做什么？

给出数组matrix.indptr中存储的每行开头的索引data。(lr, ri)每行中非零值的数据索引范围也是 如此。matrix.data[le:ri]给出该行的非零值。将其传递给您将为您提供本地索引，该索引将从后面开始np.argpartition(..., -n_row_pick)对行进行最大元素的排序。选择那些本地索引。然后将这些本地索引移回数据数组中的索引。最后，将其传递给以获得矩阵空间中最大的 n 值索引。n_row_pick[-n_row_pick:]le +matrix.indices