10
\$\begingroup\$

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)
\$\endgroup\$
2
\$\begingroup\$

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.

\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.