Search code examples
pythonlistmatrixgriddiagonal

How do I get the diagonal elements of a matrix/grid from a specific element?


I have an 8x8 grid of different numbers, and I want to get the elements of the diagonal that contains a given starting position. Here is an example

l = [[str(randint(1,9)) for i in range(8)] for n in range(8)]

>> [
[1 5 2 8 6 9 6 8]
[2 2 2 2 8 2 2 1]
[9 5 9 6 8 2 7 2]
[2 8 8 6 4 1 8 1]
[2 5 5 5 4 4 7 9]
[3 9 8 8 9 4 1 1]
[8 9 2 4 2 8 4 3]
[4 4 7 8 7 5 3 6]
]

How would I go about getting the diagonal from the position x=4 and y=3 (so 4th list and 5th element in that list)? So The diagonal I would want would be [5,2,6,4,4,1,3].


Solution

  • You can calculate the row and column of the top-left item of the diagonal based on the difference of x and y, and the number of iterations based on the difference between the lower of the two boundaries and the higher of the starting row and column:

    def diagonal(m, x, y):
        row = max((y - x, 0))
        col = max((x - y, 0))
        for i in range(min((len(m), len(m[0]))) - max((row, col))):
            yield m[row + i][col + i]
    

    so that:

    m = [
        [1, 5, 2, 8, 6, 9, 6, 8],
        [2, 2, 2, 2, 8, 2, 2, 1],
        [9, 5, 9, 6, 8, 2, 7, 2],
        [2, 8, 8, 6, 4, 1, 8, 1],
        [2, 5, 5, 5, 4, 4, 7, 9],
        [3, 9, 8, 8, 9, 4, 1, 1],
        [8, 9, 2, 4, 2, 8, 4, 3],
        [4, 4, 7, 8, 7, 5, 3, 6],
    ]
    print(list(diagonal(m, 4, 3)))
    

    outputs:

    [5, 2, 6, 4, 4, 1, 3]