有没有办法有效地将对角线的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将是一个非常大的矩阵，所以我不想单独遍历元素。出于同样的原因，我也不想预先以正确的顺序准备所有索引的列表。因为A很大，结果也同样很大，所以我想避免另外使用大量内存的解决方案。如果可以指定我想展平哪个对角线子集，那就更好了，例如，仅展平第一个和第二个对角线[1,5,9,2,6]。

Answer 1

以下功能基于indices比较，基于每个对角线都有索引关系，例如在主对角线处i==j，依此类推...

即使对非平方的2D数组也有效。

def flatten_diagonally(x, diags=None):

    diags = np.array(diags)

    if x.shape[1] > x.shape[0]:

        diags += x.shape[1]-x.shape[0]

    n = max(x.shape)

    ndiags = 2*n-1

    i,j = np.indices(x.shape)

    d = np.array([])

    for ndi in range(ndiags):

        if diags != None:

            if not ndi in diags:

                continue

        d = np.concatenate((d,x[i==j+(n-1)-ndi]))

    return d

例子：

print flatten_diagonally(A)

#[ 7.  4.  8.  1.  5.  9.  2.  6.  3.]

print flatten_diagonally(A, diags=(1,2))

#[ 4.  8.  1.  5.  9.]

对于非平方数组：

A=np.array([[1,2,3],

            [7,8,9]])

print flatten_diagonally(A, diags=(1,2))

#[ 1.  8.  2.  9.]

Answer 2

这应该为您提供所需的输出：（请注意，第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，但我不确定哪个效果更好。