python - 对大 numpy 矩阵的单元格进行排序

Question

我创建了一个大的(比如 4000 X 4000)numpy 浮点矩阵。我正在按浮点值对矩阵的单元格进行排序，生成 (row,col,value) 元组列表。这是我的代码(简化):

def cells(matrix):
  shape = np.shape(matrix)
  for row in range(shape[0]):
    for col in range(shape[1]):
      yield (row, col, matrix[row,col])

# create a random matrix
matrix = np.random.randint(100, size=(4000,4000))
# sort the cells by value
sorted_cells = sorted(cells(matrix), key=lambda x: x[2])

我知道逐个单元格的产量效率很低，但我不知道如何使用纯 numpy 迭代矩阵的 (row, col, value) 元组？ 也许这才是真正的问题!

我当前方法的问题是我的计算机在排序步骤中完全死机了。

如果我这样做，这不是问题:sorted(matrix.flatten()) 它工作正常，实际上相当快，但是我没有得到行和列...

Answer 1

numpy.argsort是你的 friend 。它不是实际对给定的数组进行排序，而是返回一个整数索引数组，告诉您如何将数组重新排序为排序顺序。鉴于此，您可以对行值和列值应用相同的排序。

这是一些代码:首先我们生成一个矩阵；在这里，我使用了不同数量的行和列，以便我们可以轻松检查结果是否正确。

>>> import numpy as np
>>> matrix = np.random.randint(100, size=(4000, 5000))
>>> rows, cols = np.indices(matrix.shape)

现在使用 argsort 获取索引。

>>> reindex = np.argsort(matrix.flatten())

使用这些索引，我们可以恢复排序后的矩阵:

>>> matrix.flat[reindex]
array([ 0,  0,  0, ..., 99, 99, 99])

和还有相应的行和列。

>>> rows.flat[reindex]
array([2455, 2870, 1196, ...,   56,   56, 3618])
>>> cols.flat[reindex]
array([ 863, 1091, 4966, ..., 3959, 3887, 4833])

为了验证答案，让我们检查第一行列对确实对应于 0 的矩阵条目，最后一行列对对应于 99:

>>> r = rows.flat[reindex]
>>> c = cols.flat[reindex]
>>> matrix[r[0], c[0]]
0
>>> matrix[r[-1], c[-1]]
99

编辑:正如 nye17 的回答所指出的，可以更直接地从 reindex 数组中恢复行和列。

>>> r, c = divmod(reindex, matrix.shape[1])

这一切都运行得非常快(排序步骤只需几秒钟)。我猜想您的原始代码锁定机器的原因是您生成的列表占用了很多内存；通过坚持使用 numpy 数组而不是列表和元组，您的内存开销会显着下降。