在计算沿轴的argmax之后，取稀疏矩阵沿该轴的最大值。

Question

我想在scipy.sparse矩阵X的一个轴上同时获取argmax和max。

>>> type(X)
scipy.sparse.csr.csr_matrix

>>> idx = X.argmax(axis=0)

>>> maxes = X.max(axis=0)

我不想计算最大值两次，但我不能像X是np.ndarray一样使用相同的方法。我该如何将argmax的索引应用于X？

- swhat

---

我可以想象复制稀疏argmax的底层代码以返回索引和值。但除此之外，我怀疑这种双重评估会是最快的。你不能简单地将密集数组的直觉转移到稀疏数组上。 - hpaulj

@hpaulj 很遗憾听到这个消息，我希望我能有稀疏矩阵的直觉... 我想知道我深入源代码所需的时间是否会比我节省的时间少。 - swhat

Csr argmax 使用 indptr 在矩阵的行上进行迭代，然后找到该行的最大值。这在概念上很简单，但有可能该行全为0，或者0本身是最小值或最大值。使用 lil 格式时，可更容易地将矩阵的行可视化。 - hpaulj

Answer 1

我通过改进你提供的方法，成功地实现了你想要的结果：

from scipy.sparse import csr_matrix

a = [[4, 0, 0], [0, 3, 0], [0, 0, 1]]
a = csr_matrix(a)
idx = a.argmax(axis=0)
m = a.shape[1]
a[idx,np.arange(m)[None,:]].toarray()

输出：

array([[4, 3, 1]], dtype=int32)

- Hemerson Tacon

---

你的 idx 超出了不同的轴。无论如何，这实际上会慢一些。稀疏矩阵索引不像密集等价物那样快速。 - hpaulj

现在正在沿所需轴提取最大值。我没有使用真正大的矩阵测试解决方案以查看您提到的速度差异是否显著。现在我在考虑是否转换为np.array然后执行argmax和max会更快。 - Hemerson Tacon

@H.Tacon，这个问题在于密集版本太大了，无法放入内存中，这也是使用稀疏矩阵的初衷。 - swhat