Pythonic way to get both diagonals passing through a matrix entry (i,j)

Question

What is the Pythonic way to get a list of diagonal elements in a matrix passing through entry (i,j)?

For e.g., given a matrix like:

[1  2  3   4  5]
[6  7  8   9 10]
[11 12 13 14 15]
[16 17 18 19 20]
[21 22 23 24 25]

and an entry, say, (1,3) (representing element 9) how can I get the elements in the diagonals passing through 9 in a Pythonic way? Basically, [3,9,15] and [5,9,13,17,21] both.

Answer 1

Using np.diagonal with a little offset logic.

import numpy as np

lst = np.array([[1,  2,  3,   4,  5],
                [6,  7,  8,   9, 10],
                [11, 12, 13, 14, 15],
                [16, 17, 18, 19, 20],
                [21, 22, 23, 24, 25]])


i, j = 1, 3
major = np.diagonal(lst, offset=(j - i))
print(major)
array([ 3,  9, 15])

minor = np.diagonal(np.rot90(lst), offset=-lst.shape[1] + (j + i) + 1)
print(minor)
array([ 5,  9, 13, 17, 21])

The indices i and j are the row and column. By specifying the offset, numpy knows from where to begin selecting elements for the diagonal.

For the major diagonal, You want to start collecting from 3 in the first row. So you need to take the current column index and subtract it by the current row index, to figure out the correct column index at the 0^th row. Similarly for the minor diagonal, where the array is flipped (rotated by 90?) and the process repeats.

Answer 2

As another alternative method, with raveling the array and for matrix with shape (n*n):

array = np.array([[1,  2,  3,   4,  5],
                  [6,  7,  8,   9, 10],
                  [11, 12, 13, 14, 15],
                  [16, 17, 18, 19, 20],
                  [21, 22, 23, 24, 25]])

x, y = 1, 3


a_mod = array.ravel()
size = array.shape[0]

if y >= x:
    diag = a_mod[y-x:(x+size-y)*size:size+1]
else:
    diag = a_mod[(x-y)*size::size+1]

if x-(size-1-y) >= 0:
    reverse_diag = array[:, ::-1].ravel()[(x-(size-1-y))*size::size+1]
else:
    reverse_diag = a_mod[x:x*size+1:size-1]

# diag         --> [ 3  9 15]
# reverse_diag --> [ 5  9 13 17 21]

The correctness of the resulted arrays must be checked further. This can be developed to handle matrices with other shapes e.g. (n*m).