Print matrix diagonals method

Question

I need to print (or store) all diagonals of a matrix. This includes the main diagonal ( { (i,i) | i = 1, 2, ... min(m,n) } ). Then all super and sub diagonals.

For example:

 1 2 3 4
 5 6 7 8
 9 10 11 12
13 14 15 16

The diagonals should be (order doesn't matter):

[
 [1, 6, 11, 16],
 [2, 7, 12],
 [5, 10, 15],
 [3, 8],
 [9, 14],
 [4],
 [13]
]

I believe I have a decent solution, but if there's a simpler / more efficient way to do this let me know. Input is assumed to be a list of lists as a NumPy array already has some methods to do this.

def print_diags(Matrix):
 m, n = len(Matrix), len(Matrix[0])
 for offset in range(min(m,n)):
 diag_upper = [row[i + offset]
 for i, row in enumerate(Matrix)
 if 0 <= i + offset < n]
 if offset != 0:
 diag_lower = [row[i - offset]
 for i, row in enumerate(matrix)
 if 0 <= i - offset < m]
 print(diag_upper, diag_lower)
 else:
 print(diag_upper)

Answer 1

Here's a generator in a generator:

m = [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]]
def diags(mat):
 width, height = len(mat[0]), len(mat)
 def diag(sx, sy):
 for x, y in zip(range(sx, height), range(sy, width)):
 yield mat[x][y]
 for sx in range(height):
 yield list(diag(sx, 0))
 for sy in range(1, width):
 yield list(diag(0, sy))

Usage:

>>> list(diags(m))
[[1, 6, 11, 16], [5, 10, 15], [9, 14], [13], [2, 7, 12], [3, 8], [4]]

---

diag starts from a coordinate and walks down the diagonal, yielding elements from it. It doesn't raise an IndexError, because zip exits after the shortest iterable is exhausted.

The last four lines are a bit inelegant, but I don't know a better way.