Sum of diagonal elements in a matrix [closed]

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.

• Your example does not work. sum_of_matrix([[1,2,3],[4,5,6],[7,8,9]]) aborts with IndexError. Jul 27, 2018 at 20:22

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

• Just for fun, [size-i-1] can be changed to [-i-1]. I'm genuinely not sure which is clearer though. Jul 27, 2018 at 21:01
• Im not sure if I can update my code, but i fixed it a little bit on my own and it works(guess I made some typo or something not sure). But your solution is definitely more efficient, so I appreciate it ! Found the typo, variables(counter and counter2) have to increment and decrement by 1, my bad! Jul 28, 2018 at 0:07
• Why special case for size == 1? It works perfectly fine without it. Jul 28, 2018 at 13:30
• Totally agree @MathiasEttinger, just thought it might be good to be aware of the possibility. Jul 28, 2018 at 23:05
• @Josiah The advantage of [-i-1] is that you can do for i, row in enumerate(matrix) and not bother with size at all. Effectively being able to reduce the function body to sum(row[i] + row[-i-1] for i, row in enumerate(matrix)). Jul 29, 2018 at 8:11