Sum of diagonal elements in a matrix [closed]

Question

The idea is to calculate sum of diagonals example [[1,2,3],[4,5,6],[7,8,9] the correct answer would be [1,5,9][3,5,7] = total 30

def sum_of_matrix(data):
    arr_solver = []
    counter = 0
    counter2 = -1
    while counter < len(data):
        arr_solver.append(data[counter][counter])
        arr_solver.append(data[counter][counter2])
        counter += data[counter][counter2]
        counter2 -= data[counter][counter]
    return sum(arr_solver)

This is my todays interview question I had, is this a good solution to a question? My idea was to implement a graph and calculate path, but that'd take waaaaay too long and probably wouldn't be able to implement it on the go.

Answer 1

The algorithm that you have implemented right now just raises an IndexError.

I'll suggest a slightly different way to do is that will have O(n) speed and O(1) memory, where n is the side length of the matrix.

Correct me if I'm wrong, but I think that's the fastest you can get.

def sum_of_diags(matrix):
    # Perhaps add some type checking first,
    # check whether the matrix is empty

    size = len(matrix[0])
    if size == 1:
        # What do you want to do with a single-element matrix?
        return matrix[0][0]*2

    # Just initializing the sum and adding to it
    # reduces the space complexity from O(n) to O(1)
    diag_sum = 0

    for i in range(size):
        # First, we sum over the main diagonal
        # from [0, 0] to [size, size]
        diag_sum += matrix[i][i]

        # Second, we sum over the other diagonal,
        # going from [0, size] to [size, 0]
        diag_sum += matrix[i][size-i-1]
    return diag_sum

Test it:

>> m = np.arange(1, 9).reshape((3, 3))
>> m
array([[1, 2, 3],
   [4, 5, 6],
   [7, 8, 9]])

>> sum_of_diags(m)
30