Получить все диагонали (включая второстепенные) многомерного массива с помощью numpy

Если вы хотите диагональ от corner1 до corner2 и определить углы в виде

(0,0,0,...,0) , (0,0,0,....,1),...,(1,1,1,...,1)

где 0 означает "Этот размер на 0", а 1 означает "Этот размер на -1/конец"

то это вернет значения, которые вы получите, перейдя от corner1 к corner2, предполагая, что массив имеет одинаковый размер в каждом измерении.

import numpy
def diagonal(arr,corner1,corner2):
    arr=numpy.array(arr)
    #Change values to fit array
    corner1Copy=(len(arr)-1)*numpy.array(corner1)
    corner2Copy=(len(arr)-1)*numpy.array(corner2)

    #create return array by running from corner1 to corner2 and returning the values
    return [arr[tuple((i*corner2Copy+(len(arr)-i-1)*corner1Copy)/(len(arr)-1))] for i in range(len(arr))]

Вот два небольших тестовых случая, но я бы предложил создать еще несколько, если я что-то упустил:

arr=[[[i+j+k for i in range(5)]for j in range(5)] for k in range(5)]
corner1=[0,0,0]
corner2=[1,1,1]

#returns arr[0,0,0],arr[1,1,1],....,arr[-1,-1,-1]
print(diagonal(arr,corner1,corner2))
print([arr[i][i][i] for i in range(len(arr))])

arr2=[[i+j for i in range(5)]for j in range(5)]

corner12=[0,1]
corner22=[1,1]
#return arr[0,-1],arr[1,-1],....,arr[-1,-1]
print(diagonal(arr2,corner12,corner22))
print([arr2[i][-1] for i in range(len(arr2))])

Это должно сделать это, используя массивы numpy. Это дает генератор со всеми диагоналями, которые существуют в том же количестве измерений, что и массив. Также дает представление (а не копию) исходного массива. Объяснение находится под кодом.

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.