Diagonal in matrix - python

Question

I need to write a function that will return the sum of numbers that form a diagonal in any given matrix.
Being a Python newbie I've got a problem. This is my code:

def diagonal(matrix):
    return sum([matrix[i][i] for i in range(len(matrix))])

I've been trying for some time now but don't see what is wrong because it always gives me back report about error saying "list index out of range".

I am not allowed to import numpy.

Any form of help, hint, advice would be appreciated.

Answer 1

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.

Answer 2

You have to stop when either of indices exceds respective dimension, for example you can limit matrix with slicing:

def diagonal_sum(matrix):
    row_size = len(matrix[0])
    return sum(row[i] for i, row in enumerate(matrix[:row_size]))

Demo:

>>> diagonal_sum([[1,2],[3,4],[5,6]])
5

Answer 3

I think diagonal is not defined for non-square matrices. So we'd better not choose the min of two dimensions only to let the code return something.

So how about instead:

def diagonal(matrix):
    try:
        return sum([matrix[i][i] for i in range(len(matrix))])
    except IndexError:
        print "Bad Matrix! :("