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

Question

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 :

out put:

[1,1,0,1,0]

2) right diagonal line:

out put:

[0,0,1,2,0]

Is there an easy way? Thanks~

Answer 1

Approach #1 : Using masking

Here's a vectorized one based on masking -

def left_diag_sum_masking(a):
    n = len(a)
    N = 2*n-1

    R = np.arange(N)
    r = np.arange(n)

    mask = (r[:,None] <= R) & (r[:,None]+n > R)

    b = np.zeros(mask.shape,dtype=a.dtype)
    b[mask] = a.ravel()

    return b.sum(0)

So, left_diag_sum_masking gives us left-diagonal summations. To get the right-diagonal ones, simply flip along cols, sum and then flip it back.

Hence, simply do -

right_diag_sum = left_diag_sum_masking(a[::-1])[::-1]

Sample run -

In [220]: np.random.seed(0)

In [221]: a = np.random.randint(0,9,(4,4))

In [222]: a
Out[222]: 
array([[5, 0, 3, 3],
       [7, 3, 5, 2],
       [4, 7, 6, 8],
       [8, 1, 6, 7]])

In [223]: left_diag_sum_masking(a)
Out[223]: array([ 5,  7, 10, 23,  9, 14,  7])

In [224]: left_diag_sum_masking(a[::-1])[::-1] # right-diag sums
Out[224]: array([ 3,  5, 13, 21, 20,  5,  8])

Approach #2 : Using zeros-padding

def left_diag_sum_zerospad(a):
    n = len(a)
    N = 2*n-1
    p = np.zeros((n,n),dtype=a.dtype)
    ap = np.concatenate((a,p),axis=1)
    return ap.ravel()[:n*N].reshape(n,-1).sum(0)

To get right-diagonal summations -

right_diag_sum = left_diag_sum_zerospad(a[::-1])[::-1]