python - 有没有办法有效地将对角线的numpy数组展平？

我正在寻找一种有效的方法（最好是一个矢量化的快速内置函数）来按对角线顺序展平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将是一个非常大的矩阵，所以我不想单独迭代元素。出于同样的原因，我也不想事先按正确的顺序准备所有索引的列表。因为A很大，结果也会同样大，所以我希望避免使用大量内存的解决方案。
如果我能指定要展平的对角线子集（例如，仅展平第一个和第二个对角线）会更好。

最佳答案

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.diag而不是np.diag，我不太确定哪个更好。Documentation

关于python - 有没有办法有效地将对角线的numpy数组展平？，我们在Stack Overflow上找到一个类似的问题：https://stackoverflow.com/questions/16866395/