У меня есть трехмерный массив вроде

a = np.array(
    [[[1, 2],
      [3, 4]],

     [[5, 6],
      [7, 8]]]
)

И я хочу получить диагональный компонент из этого массива, я имею в виду,

>> np.array([a[i,i,:] for i in range(min(a.shape[0], a.shape[1]))])
array([[1, 2],
       [7, 8]])

Есть ли более быстрый способ сделать то же самое, что и выше?

0
seekworser 3 Янв 2018 в 08:10

2 ответа

Лучший ответ

Вместо np.diagonal вы также можете использовать np.einsum:

>>> a = np.arange(1,9).reshape(2,2,2)
>>> d = np.einsum('iij->ij', a)
>>> d
array([[1, 2],
       [7, 8]])

Преимущество: он возвращает доступное для записи представление, поэтому его можно использовать для непосредственного управления диагональю в исходном массиве:

>>> d[0,0] = -3
>>> a
array([[[-3,  2],
        [ 3,  4]],

       [[ 5,  6],
        [ 7,  8]]])
1
Paul Panzer 3 Янв 2018 в 07:46

np.diagonal работает для массивов с более чем двумя измерениями, которые вы можете указать ось для просмотра с параметрами axis1 и axis2, для вашего случая работает значение по умолчанию (axis1 = 0, axis2 = 1):

np.diagonal(a, axis1=0, axis2=1).T
# array([[1, 2],
#        [7, 8]])

Это также работает для массивов, которые имеют разные размеры по двум рассматриваемым осям:

a = np.array(
    [[[1, 2],
      [3, 4]],
​
     [[5, 6],
      [7, 8]],

     [[9, 10],
      [11, 12]]]
)
​
np.diagonal(a).T
#array([[1, 2],
#       [7, 8]])

a = np.array(
    [[[1, 2],
      [3, 4],
      [9, 10]],
​
     [[5, 6],
      [7, 8],
      [11, 12]]]
)
​
np.diagonal(a).T
#array([[1, 2],
#       [7, 8]])

Или вы можете использовать расширенное индексирование :

Создайте индекс диапазона:

idx = np.arange(min(a.shape[:2]))

idx
# array([0, 1])

Проиндексируйте первое и второе измерение с помощью целочисленных массивов, а третье измерение с помощью среза:

a[idx, idx, :]
#array([[1, 2],
#       [7, 8]])
1
Psidom 3 Янв 2018 в 05:42