Search code examples
pythonpython-3.xnumpydiagonal

Sum along diagonal and anti-diagonal lines in 2D array - NumPy / Python

I have a np array like below:

np.array([[1,0,0],[1,0,0],[0,1,0]])
output:
array([[1, 0, 0],
       [1, 0, 0],
       [0, 1, 0]])

I wish to sum left diagonal and right right diagonal line elements to new array:

1) left diagonal line :

enter image description here

out put:

[1,1,0,1,0]

2) right diagonal line:

enter image description here

out put:

[0,0,1,2,0]

Is there an easy way? Thanks~

Solution

  • you can use .diag like this:

    import numpy as np

arr = np.array([[1, 0, 0], [1, 0, 0], [0, 1, 0]])

left_diag = [np.sum(np.diag(np.fliplr(arr), d)) for d in range(len(arr) - 1, -len(arr), -1)]
right_diag = [np.sum(np.diag(arr, d)) for d in range(len(arr) - 1, -len(arr), -1)]

print("left:", left_diag)
print("right:", right_diag)

    Output:

    left: [1, 1, 0, 1, 0]
right: [0, 0, 1, 2, 0]

    it returns all the element in a diagonal of a given offset. to get all the offsets in the order you mentioned, we go from +2 to -2, then for each diagonal get the sum of elements.

    to get the left diagonal we first flip arr using .fliplr