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):
        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.

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

1 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)
  • 2
    \$\begingroup\$ Just for fun, [size-i-1] can be changed to [-i-1]. I'm genuinely not sure which is clearer though. \$\endgroup\$
    – Josiah
    Jul 27, 2018 at 21:01
  • \$\begingroup\$ 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! \$\endgroup\$
    – nexla
    Jul 28, 2018 at 0:07
  • 1
    \$\begingroup\$ Why special case for size == 1? It works perfectly fine without it. \$\endgroup\$ Jul 28, 2018 at 13:30
  • \$\begingroup\$ Totally agree @MathiasEttinger, just thought it might be good to be aware of the possibility. \$\endgroup\$ Jul 28, 2018 at 23:05
  • 1
    \$\begingroup\$ @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)). \$\endgroup\$ Jul 29, 2018 at 8:11

Not the answer you're looking for? Browse other questions tagged or ask your own question.