Диагональ в NumPy

Если у квадратной таблицы обе оси обозначены одной и той же буквой на входе, индексы строки и столбца совпадают только на главной диагонали. Одна буква справа оставляет ряд из этих элементов, без суммирования.

Запись 'ii->i' повторяет 'i' у входной таблицы и возвращает главную диагональ, как при знакомом diag для двумерного массива:

import numpy as np

square = np.array([[10, 1],
                   [2, 20]])
main = np.einsum('ii->i', square)
print(main) # выведет [10 20]

Для таблицы [[5, 9], [7, 8]] получите главную диагональ формулой с повторяющейся буквой на входе и одной буквой в ответе и выведите ряд.

Для таблицы баллов [[0, 1], [2, 3]] извлеките элементы, где номер строки равен номеру столбца, и выведите результат.