在 Python 中获取矩阵/列表列表中的所有对角线 [英] Get all the diagonals in a matrix/list of lists in Python

问题描述

我正在寻找一种 Pythonic 方法来获取(方形)矩阵的所有对角线，表示为列表列表.

I'm looking for a Pythonic way to get all the diagonals of a (square) matrix, represented as a list of lists.

假设我有以下矩阵:

matrix = [[-2,  5,  3,  2],
          [ 9, -6,  5,  1],
          [ 3,  2,  7,  3],
          [-1,  8, -4,  8]]

那么大对角线很容易:

l = len(matrix[0])
print([matrix[i][i] for i in range(l)])              # [-2, -6, 7,  8]
print([matrix[l-1-i][i] for i in range(l-1,-1,-1)])  # [ 2,  5, 2, -1]

但是我在想出一种生成所有对角线的方法时遇到了麻烦.我正在寻找的输出是:

But I have trouble coming up with a way to generate all the diagonals. The output I'm looking for is:

[[-2], [9, 5], [3,-6, 3], [-1, 2, 5, 2], [8, 7, 1], [-4, 3], [8],
 [2], [3,1], [5, 5, 3], [-2, -6, 7, 8], [9, 2, -4], [3, 8], [-1]]

推荐答案

在 numpy<中可能有更好的方法来做到这一点/a> 比下面，但我还不太熟悉它:

There are probably better ways to do it in numpy than below, but I'm not too familiar with it yet:

import numpy as np

matrix = np.array(
         [[-2,  5,  3,  2],
          [ 9, -6,  5,  1],
          [ 3,  2,  7,  3],
          [-1,  8, -4,  8]])

diags = [matrix[::-1,:].diagonal(i) for i in range(-3,4)]
diags.extend(matrix.diagonal(i) for i in range(3,-4,-1))
print [n.tolist() for n in diags]

输出

[[-2], [9, 5], [3, -6, 3], [-1, 2, 5, 2], [8, 7, 1], [-4, 3], [8], [2], [3, 1], [5, 5, 3], [-2, -6, 7, 8], [9, 2, -4], [3, 8], [-1]]

<小时>

编辑:已更新以适用于任何矩阵大小.

---

Edit: Updated to generalize for any matrix size.

import numpy as np

# Alter dimensions as needed
x,y = 3,4

# create a default array of specified dimensions
a = np.arange(x*y).reshape(x,y)
print a
print

# a.diagonal returns the top-left-to-lower-right diagonal "i"
# according to this diagram:
#
#  0  1  2  3  4 ...
# -1  0  1  2  3
# -2 -1  0  1  2
# -3 -2 -1  0  1
#  :
#
# You wanted lower-left-to-upper-right and upper-left-to-lower-right diagonals.
#
# The syntax a[slice,slice] returns a new array with elements from the sliced ranges,
# where "slice" is Python's [start[:stop[:step]] format.

# "::-1" returns the rows in reverse. ":" returns the columns as is,
# effectively vertically mirroring the original array so the wanted diagonals are
# lower-right-to-uppper-left.
#
# Then a list comprehension is used to collect all the diagonals.  The range
# is -x+1 to y (exclusive of y), so for a matrix like the example above
# (x,y) = (4,5) = -3 to 4.
diags = [a[::-1,:].diagonal(i) for i in range(-a.shape[0]+1,a.shape[1])]

# Now back to the original array to get the upper-left-to-lower-right diagonals,
# starting from the right, so the range needed for shape (x,y) was y-1 to -x+1 descending.
diags.extend(a.diagonal(i) for i in range(a.shape[1]-1,-a.shape[0],-1))

# Another list comp to convert back to Python lists from numpy arrays,
# so it prints what you requested.
print [n.tolist() for n in diags]

输出

[[ 0  1  2  3]
 [ 4  5  6  7]
 [ 8  9 10 11]]

[[0], [4, 1], [8, 5, 2], [9, 6, 3], [10, 7], [11], [3], [2, 7], [1, 6, 11], [0, 5, 10], [4, 9], [8]]

这篇关于在 Python 中获取矩阵/列表列表中的所有对角线的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！