Python 如何有效地在多个矩阵中分别为每个元素找到N个最大值？

我正在遍历大量的hxw矩阵。我无法将它们全部存储在内存中。我需要得到N个矩阵。例如，位于（i，j）位置的N个矩阵中的第1个矩阵的元素将是所有已处理矩阵中位于（i，j）位置的所有元素中最大的元素。对于N矩阵的第二个，将采用第二大的元素，依此类推

例如

设N=2。那么第一个矩阵将如下所示

第二个矩阵是这样的

---

如何在循环内执行这样的操作，以避免将所有矩阵存储在内存中？

建议使用

np.partition

函数进行注释。我将numpy的使用替换为，后者使用GPU。并且还添加了一个缓冲区来减少排序的频率

import cupy as np

buf = // # As much as fits into the GPU
largests = np.zeros((buf + N, h, w))
for i in range(num):
    val = //
    largests[i % buf] = val
    if i % buf == buf - 1:
        largests.partition(range(buf, buf + N), axis=0)
largests.partition(range(buf, buf + N), axis=0)  # Let's not forget the tail
res = largests[:-(N + 1):-1]

---

解决方案的速度不是很快，但我已经适应了这种速度。

你想对矩阵进行排序吗？@Kevin，我不明白你的问题。这个问题可以通过对所有矩阵排序和切片来解决。但问题是所有矩阵都必须存储在内存中。与排序相比，在循环的一个过程中可以找到N=const最大值。看起来是

hxw

@QuangHoang，是的，你是对的。我甚至可以尝试实现一个简单的算法。但是numpy允许你方便而有效地对操作进行矢量化，我想知道一个简单而有效的解决方案，但是这也需要把所有的东西都放在内存中。我认为

np

在流数据方面没有任何优势。