У меня есть трехмерный массив вроде
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]])
Есть ли более быстрый способ сделать то же самое, что и выше?
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]]])
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]])
Похожие вопросы
Новые вопросы
python-3.x
По вопросам программирования Python, относящегося к версии языка 3+. Используйте более общий тег [python] для всех вопросов Python и добавляйте его только в том случае, если ваш вопрос зависит от версии. Используйте теги [python-2.x] для вопросов Python 2.
En Español