Написание диагональной функции матрицы (2D списка)

Question

Я хочу написать функцию, которая принимает один аргумент, квадратную 2D-матрицу и возвращает диагональ введенной матрицы из интерактивного окна IDLE.

Я пробовал несколько способов, но всякий раз, когда я вводил матрицу (от 1 до 16 в 4 & times; 4), я получал TypeError, говорящий 4 positional arguments but 5 were given.

Я возился с этим, пытаясь использовать input(), чтобы иметь возможность ввести желаемую матрицу, но я все еще получаю ту же ошибку.

Вот что я хочу сделать:

>>> m1 = [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]]
>>> diagonal(m1)
[1, 6, 11, 16]

Как бы я поступил об этом? Я также видел много постов, использующих Numpy, но я не совсем уверен, разрешено ли использовать это назначение.

Answer 1

enumerate() очень удобно здесь с квадратной матрицей, потому что она даст вам правильные индексы при циклическом просмотре массивов.Вы можете использовать его в понимании, чтобы получить и правильный индекс, и правильную матрицу.

def diagonal(mat):
    return [el[i] for i, el in enumerate(mat)]

m1 = [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]]

diagonal(m1)
>> [1, 6, 11, 16]

Answer 2

Попробуйте это:

def diagonal(m1):
    n = len(m1)
    l = []
    for i in range(n):
        l.append(m1[i][i])

    return l

m1 = [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]]
print(diagonal(m1))