pythonmatrixdiagonal

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


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

Solution

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