Computing diagonal sums in numpy

user_16715836 0
Viewed 159

I have a fairly large rectangular numpy array, of shape (m, n), for example:

>>> a.shape
(27584, 34092)

I have to compute the sum of each anti-diagonal of the array. This new array will have a shape of (m + n - 1,).

The simplistic approach is to do:

m, n = a.shape
r = np.zeros(m + n - 1)
for i in range(m):
    for j in range(n):
        r[i + j] += a[i][j]

# r is the sum of all anti-diagonals of a

This is obviously very slow, is there any way to perform the calculation using a clever numpy primitive? My only other option would be to code this in C++, which is also doable too - but requires more work.

array-broadcasting numpy-ndarray numpy python
1 Answers

In my opinion, the most intuitive solution is to:

  • use a[::-1] to "flip" the order of rows, so that then you can operate on "normal" diagonals,
  • use np.diagonal to access each diagonal,
  • sum the returned diagonal,
  • put the above code in a list comprehension.

So the whole code can be:

result = [np.diagonal(a[::-1], offs).sum() for offs in range(-a.shape[0] + 1, a.shape[1])]

It should run significantly faster than your code.

Since your source array is quite a big one, maybe you will get some speed gain by:

  • computing the "flipped" array once, at the beginning,
  • then operate on this array:

So the reworked code, including a hint about np.trace (see comments), can be:

b = a[::-1]
result = [np.trace(b, offs) for offs in range(-a.shape[0] + 1, a.shape[1])]
user_73884770
Related