Элементы главной диагонали

Задание

Дан набор данных об оценках студентов по предметам.

Оставить только данные, находящиеся на главной диагонали.

Гарантируется, что входные данные с оценками будут представлены в виде квадратной матрицы.

Например

import pandas as pd

data = {
    "name": ["Иван", "Юлия", "Андрей", "Наталья"],
     "subject_1": [1, 5, 2, 4],
     "subject_2": [2, 5, 3, 1],
     "subject_3": [1, 5, 2, 5],
     "subject_4": [3, 5, 1, 2]
}

df = pd.DataFrame(data)

#       name  subject_1  subject_2  subject_3  subject_4
# 0     Иван          1          2          1          3
# 1     Юлия          5          5          5          5
# 2   Андрей          2          3          2          1
# 3  Наталья          4          1          5          2

# Правильный ответ
# [1, 5, 2, 2]