Main Diagonal Extraction

Problem

Return the main diagonal of matrix A as a list: d[k] = A[k][k].

A may be rectangular; stop when either index runs out.

Hint

The diagonal length is min(rows, columns).

case 1

in: [[[1,2],[3,4]]]

expected: [1,4]

case 2

in: [[[5,0,0],[0,6,0],[0,0,7]]]

expected: [5,6,7]

case 3

in: [[[1,2,3],[4,5,6]]]

expected: [1,5]

case 4

in: [[[9]]]

expected: [9]

Starter codePython

def extract_diagonal(A):
    # Your code here
    pass

4 test cases · run your code against them in the editor