get the sum of diagonals and anti diagonals 2D N*N matrix using numpy

Question

Consider an N*N matrix. I select any element and then I want the sum of diagonals and antidiagonals passing through it using numpy.

For example:

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

For the number 8 it should return (2+8+14) + (4+8+12+16+20).

Thanks.

Answer 1

Use np.diag. Suppose you select element a[i, j]. Then,

>>> diagonal = np.diag(a, j-i)
>>> antidiagonal = np.diag(a[:, ::-1], N-j-1-i)
>>> np.sum(diagonal) + np.sum(antidiagonal)

This works, because j-i is the number of columns above the diagonal entry (i, i). So np.diag(a, j-i) correctly picks the diagonal corresponding to 8.

For the antidiagonal, I simply flip the array horizontally (by reversing the columns), and then run np.diag again. But now, the new column index is N-j-1, while the row index remains i.

If you find a[:, ::-1] unreadable, you could use np.fliplr(a) or np.flipud(a) instead, as Eric suggests.