Get all the diagonals in a matrix/list of lists in Python

Question

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

Suppose I have the following matrix:

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

Then the large diagonals are easy:

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]] 

Answer 1

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] 

Output

[[-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] 

Output

[[ 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]]