排序大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])

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

我目前的方法的问题是我的计算机在分拣步骤中完全死亡.

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

Answer 1

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

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

>>> 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数组而不是列表和元组,你的内存开销显着下降.