python - 有没有办法有效地按对角线顺序展平 numpy 数组？

Question

我正在寻找一种有效的方法(最好是向量化的快速内置函数)来按对角线顺序展平 numpy 数组。例如:

A=np.array([[1,2,3],
            [4,5,6],
            [7,8,9]])
b=flatten_diagonally(A)

b 应该是 [7,4,8,1,5,9,2,6,3]。

A 将是一个非常大的矩阵，所以我不想单独迭代元素。出于同样的原因，我也不想提前准备一份按正确顺序排列的所有指数的 list 。由于 A 很大，结果也同样大，因此我想避免另外使用大量内存的解决方案。

如果我可以指定我想要展平的对角线子集，那就更好了，例如仅展平第 1 条和第 2 条对角线将得到 [1,5,9,2,6]。

Answer 1

numpy.diag 返回沿特定索引的对角线。 Documentation

所以这应该给你想要的输出:(注意第 0 条对角线是正常对角线，所以如果你想要次对角线，你可能需要对角线使用负值。)

import numpy as np

def flatten_diagonally(npA, diagonals = None):
    diagonals = diagonals or xrange(-npA.shape[0] + 1, npA.shape[1])
    return np.concatenate(map(lambda x: np.diag(npA, k = x), diagonals))

请注意，您可以使用 np.diagonal 而不是 np.diag，我不太确定哪个效果更好。 Documentation