If you are sure that your matrix is rectangular (len(matrix[i])
is the same for all lists in matrix
), then you can sum your list only as long as your smaller dimension goes:
def diagonal(matrix):
return sum([matrix[i][i] for i in range(min(len(matrix[0]),len(matrix)))])
len(matrix)
is the first dimension of your matrix, and len(matrix[0])
is the dimension of the first row vector, which is the second dimension for rectangular matrices.