Get all diagonals (including minor ones) of a multi-dimensional array with numpy

Viewed 1312

Im trying to get the diagonal (and anti-diagonal) elements of multi-dimensional objects.

The shapes are like (2,2), (3,3,3), (4,4,4,4), (5,5,5,5,5) and so on. I don't think this is too relevant though.

I found ways of getting the diagonal elements, with the .diagonal method of the ndarray, but I can't find anything that would get me the antidiagonal.

So would I have to do this by hand?

[EDIT] So for

array([[[54, 81, 31],
        [ 4, 83, 18],
        [38, 32, 52]],

       [[ 2, 45, 87],
        [33, 20,  3],
        [85, 31, 35]],

       [[ 6, 11, 49],
        [39, 76, 75],
        [28, 52, 63]]])

So I'd want the "horizontal" diagonals, like:

[54, 45, 49],
[ 4. 20, 75],
etc.

but then these are also horizontal in some sense

[ 6, 45, 31],
[39, 20, 18]

and then "vertical" ones like:

[54, 33, 28],
[81, 20, 52],
etc.

but then these are also vertical:

[6, 33, 38],
[11, 20, 32]

and then this one, however you'd call it

[54, 20, 63]

and then these are also "longer" diagonals, like the previous one (longer in a geometrical sense, if you think of the matrix as a 3d geometrical structure, with the numbers being placed on the vertexes of a cube, and on the middle of the lines between them)

[38, 20, 49],
[6, 20, 52]

Then, a minor diagonal would be one which goes from right to left or bottom to top (but not both at the same time) in this matrix, something like:

[31, 45, 6],
[31, 83, 38]  # this is the first classical anti-diagonal in the first matrix

Of course, I did'nt put here all the diagonals, but this is my requirement. I don't need diagonals which are offset from any of the main/anti diagonals.

If you also know this not to be possible, please tell, because I will do it by hand then.

2 Answers

This should do it, using numpy arrays. It gives a generator with all the diagonals that exist in the same number of dimensions as the array. Also gives a view (rather than a copy) of the original array. Explanation is below the code.

import numpy as np
def get_diagonals_np(arr):
    if arr.ndim == 1:
        yield arr
    else:
        yield from get_diagonals_np(arr.diagonal())
        yield from get_diagonals_np(np.flip(arr, 0).diagonal())

# The function is recursive. How it works is best shown by example.
# 1d: arr = [0, 1] then the diagonal is also [0, 1].

# 2d: arr = [[0, 1],
#            [2, 3]]
# The numpy diagonal method gives the main diagonal = [0, 3], a 1d array
# which is recursively passed to the function.
# To get the opposite diagonal we first use the numpy flip function to
# reverse the order of the elements along the given dimension, 0 in this case.
# This gives [[2, 3],
#              0, 1]]
# The numpy diagonal method gives the main diagonal = [2, 1], a 2d array
# which is recursively passed to the function.

# 3d: arr = [[[0, 1],
#             [2, 3]],
#            [[4, 5],
#             [6, 7]]]
# The numpy diagonal method gives the main diagonals in the 3rd dimension
# as rows.
#            [[0, 6],
#             [1, 7]]
# Note that the diagonals of this array are [0, 7] and [6, 1] which are
# retrieved by a recurive call to the function.
# We now have 2 of the 4 3-agonals of the orginal 3d arr.
# To get the opposite 3-agonals we first use the numpy flip function which
# gives
#           [[[4, 5],
#             [6, 7]],
#            [[0, 1],
#             [2, 3]]]
# and a call to the numpy diagonal method gives
#            [[4, 2],
#             [5, 3]]
# The diagonals of this array are [4, 3] and [2, 5]
# We now have all four 3-agonals of the original 3d arr.
Related