Python 3: Get diagonal elements of matrix (list of lists) about a point in that matrix without NumPy

Question

I have a matrix (formed of a list of lists) that would look something like:

matrix = [[0, 0, 0, 0, 5],
          [0, 0, 0, 4, 0],
          [2, 0, 3, 0, 0],
          [3, 2, 0, 2, 0],
          [1, 0, 2, 0, 1]]

What I am struggling to create is a function that will take this matrix as an input, along with a position in the matrix - represented by a tuple - and return the two diagonals that intersect that point (without using NumPy). For example,

def getDiagonal(matrix, pos)
    (row, col) = pos
    # Smart diagonal finder code #
    return (diag1, diag2)

diagonals = getDiagonals(matrix, (1, 1))
print(diagnonal[0])
print(diagnonal[1])

print(' ')

diagonals = getDiagonals(matrix, (1, 3))
print(diagnonal[0])
print(diagnonal[1])

Expected output:

OUT: [5, 4, 3, 2, 1]
OUT: [2, 2, 2]
OUT: 
OUT: [0, 2, 2]
OUT: [0, 0, 3, 2, 1]

It is worth pointing out that I don't mind from which direction (bottom-to-top or top-to-bottom) the returned elements of the diagonals are. They could easily be done one way and revered using reverse() if need be.

I have looked at similar questions such as this one but this mainly deals with acquring the leading diagonals of a matrix and provides less information on getting the diagonals about a point.

Many thanks for your help and comments in advance!

Answer 1

A bit confusing, but I think this does it:

def getDiagonals(matrix, pos):
    row, col = pos
    nrows = len(matrix)
    ncols = len(matrix[0]) if nrows > 0 else 0
    # First diagonal
    d1_i, d1_j = nrows - 1 - max(row - col, 0), max(col - row, 0)
    d1_len = min(d1_i + 1, ncols - d1_j)
    diag1 = [matrix[d1_i - k][d1_j + k] for k in range(d1_len)]
    # Second diagonal
    t = min(row, ncols - col - 1)
    d2_i, d2_j = nrows - 1 - row + t, col + t
    d2_len = min(d2_i, d2_j) + 1
    diag2 = [matrix[d2_i - k][d2_j - k] for k in range(d2_len)]
    return (diag1, diag2)

# Test
matrix = [[0, 0, 0, 0, 5],
          [0, 0, 0, 4, 0],
          [2, 0, 3, 0, 0],
          [3, 2, 0, 2, 0],
          [1, 0, 2, 0, 1]]
diagonals = getDiagonals(matrix, (1, 1))
print(diagonals[0])
# [1, 2, 3, 4, 5]
print(diagonals[1])
# [2, 2, 2]

diagonals = getDiagonals(matrix, (1, 3))
print(diagonals[0])
# [2, 2, 0]
print(diagonals[1])
# [1, 2, 3, 0, 0]

diagonals = getDiagonals(matrix, (2, 2))
print(diagonals[0])
# [1, 2, 3, 4, 5]
print(diagonals[1])
# [1, 2, 3, 0, 0]