Problem Description

You are given a square matrix mat (meaning it has the same number of rows and columns). Your task is to calculate the sum of all elements on both diagonals of the matrix.

The matrix has two diagonals:

• Primary diagonal: Elements where the row index equals the column index (i.e., mat[0][0], mat[1][1], mat[2][2], etc.)
• Secondary diagonal: Elements where the row index plus column index equals n-1, where n is the size of the matrix (i.e., mat[0][n-1], mat[1][n-2], mat[2][n-3], etc.)

You need to sum all elements on both diagonals. However, be careful with the center element in odd-sized matrices - if an element appears on both diagonals (which happens only at the center of odd-sized matrices), count it only once in your sum.

For example, in a 3×3 matrix:

• Primary diagonal elements are at positions [0][0], [1][1], [2][2]
• Secondary diagonal elements are at positions [0][2], [1][1], [2][0]
• The element at [1][1] is on both diagonals, so it should be counted only once

The solution iterates through each row i and adds the element at position [i][i] (primary diagonal) and the element at position [i][n-i-1] (secondary diagonal). When these two positions are the same (i.e., i == n-i-1), which happens at the center of odd-sized matrices, the element is added only once.

Intuition

The key insight is that in a square matrix, we can identify diagonal elements using a simple pattern. For any row i, the primary diagonal element is always at column i (same as the row index), and the secondary diagonal element is at column n - i - 1 (counting from the opposite end).

Instead of thinking about diagonals as separate entities to traverse, we can process the matrix row by row. As we go through each row, we know exactly where the diagonal elements are located without needing to search for them.

For a row at index i:

• Primary diagonal element: mat[i][i]
• Secondary diagonal element: mat[i][n - i - 1]

This pattern holds for every row in the matrix. The beauty of this approach is that we only need a single loop through the rows, and for each row, we can directly access both diagonal elements in constant time.

The tricky part is handling the center element in odd-sized matrices. When the matrix has an odd dimension (like 3×3 or 5×5), there's a center element that belongs to both diagonals. This happens when i == n - i - 1, which simplifies to i == (n - 1) / 2 - exactly the middle row. At this point, both diagonal positions point to the same element, so we need to add it only once to avoid double-counting.

This row-by-row approach is efficient because we visit each diagonal element exactly once, giving us a time complexity of O(n) where n is the number of rows, rather than O(n²) if we were to check every element in the matrix.

Solution Approach

The implementation follows a straightforward row-by-row traversal pattern. We iterate through each row of the matrix and calculate the sum of diagonal elements efficiently.

Here's how the algorithm works:

1. Initialize the sum: Start with ans = 0 to accumulate the diagonal sum.

2. Get matrix dimension: Store n = len(mat) to know the size of our square matrix.

3. Iterate through rows: Use enumerate(mat) to get both the row index i and the row content row simultaneously.

4. Calculate diagonal positions: For each row i:

   • The primary diagonal element is at position row[i]
   • The secondary diagonal element position is calculated as j = n - i - 1
   • The secondary diagonal element is at position row[j]

5. Handle the center element: When adding elements to the sum:

   • Always add the primary diagonal element: row[i]
   • For the secondary diagonal element, check if j == i (which means we're at the center of an odd-sized matrix)
   • If j == i, add 0 (to avoid double-counting the same element)
   • Otherwise, add row[j]

   This is elegantly handled with the expression: ans += row[i] + (0 if j == i else row[j])

6. Return the result: After processing all rows, return the accumulated sum.

The algorithm makes a single pass through the matrix rows, accessing exactly two elements per row (or one element for the center row in odd-sized matrices). This gives us:

• Time Complexity: O(n) where n is the number of rows
• Space Complexity: O(1) as we only use a constant amount of extra space

The solution elegantly avoids nested loops and unnecessary condition checks by using the mathematical relationship between row indices and diagonal positions.

Example Walkthrough

Let's walk through the solution with a 3×3 matrix:

mat = [[1, 2, 3],
       [4, 5, 6],
       [7, 8, 9]]

Step 1: Initialize

• ans = 0 (our running sum)
• n = 3 (matrix dimension)

Step 2: Process Row 0 (i=0)

• Row content: [1, 2, 3]
• Primary diagonal element: row[0] = 1
• Secondary diagonal position: j = 3 - 0 - 1 = 2
• Secondary diagonal element: row[2] = 3
• Check if center: j == i? → 2 == 0? → No
• Add to sum: ans = 0 + 1 + 3 = 4

Step 3: Process Row 1 (i=1)

• Row content: [4, 5, 6]
• Primary diagonal element: row[1] = 5
• Secondary diagonal position: j = 3 - 1 - 1 = 1
• Secondary diagonal element: row[1] = 5
• Check if center: j == i? → 1 == 1? → Yes! (This is the center)
• Add to sum: ans = 4 + 5 + 0 = 9 (add 0 instead of 5 to avoid double-counting)

Step 4: Process Row 2 (i=2)

• Row content: [7, 8, 9]
• Primary diagonal element: row[2] = 9
• Secondary diagonal position: j = 3 - 2 - 1 = 0
• Secondary diagonal element: row[0] = 7
• Check if center: j == i? → 0 == 2? → No
• Add to sum: ans = 9 + 9 + 7 = 25

Final Result: 25

The diagonal elements we summed were:

• Primary diagonal: 1, 5, 9 (sum = 15)
• Secondary diagonal: 3, 5, 7 (sum = 15)
• Center element 5 was counted only once
• Total: 1 + 5 + 9 + 3 + 7 = 25

Solution Implementation

Time and Space Complexity

The time complexity is O(n), where n is the number of rows in the matrix. The algorithm iterates through each row exactly once using the enumerate function, performing constant-time operations (array access and arithmetic) for each row. Since there are n rows and we visit each row once with O(1) work per row, the total time complexity is O(n).

The space complexity is O(1). The algorithm only uses a fixed amount of extra space regardless of the input size. The variables ans, n, i, row, and j are the only additional space used, and these do not scale with the size of the input matrix. The enumerate iterator and the reference to each row do not create new data structures but rather reference existing elements in the input matrix.

Learn more about how to find time and space complexity quickly.

Common Pitfalls

1. Double-Counting the Center Element in Odd-Sized Matrices

The Pitfall: The most common mistake is forgetting to handle the center element in odd-sized matrices. When the matrix has odd dimensions (e.g., 3×3, 5×5), the center element belongs to both diagonals. A naive implementation might add this element twice:

# Incorrect approach - double counts center element
def diagonalSum_wrong(mat):
    total = 0
    n = len(mat)
    for i in range(n):
        total += mat[i][i]  # Primary diagonal
        total += mat[i][n-i-1]  # Secondary diagonal - might double count!
    return total

For a 3×3 matrix, when i = 1, both mat[1][1] (primary) and mat[1][3-1-1] = mat[1][1] (secondary) refer to the same center element, causing it to be counted twice.

The Solution: Check if the current position on the primary diagonal is the same as the secondary diagonal position before adding:

def diagonalSum_correct(mat):
    total = 0
    n = len(mat)
    for i in range(n):
        total += mat[i][i]  # Always add primary diagonal
        if i != n - i - 1:  # Only add secondary if different from primary
            total += mat[i][n-i-1]
    return total

2. Using Incorrect Index Calculation for Secondary Diagonal

The Pitfall: Another frequent error is miscalculating the column index for the secondary diagonal. Some might incorrectly use formulas like n - i or n - 1 - i inconsistently:

# Incorrect - off by one error
secondary_col = n - i  # Should be n - i - 1

# This would try to access mat[0][n] which is out of bounds

The Solution: Remember that for a matrix of size n, the secondary diagonal element in row i is at column n - i - 1. This ensures:

• Row 0: column n-1 (last column)
• Row 1: column n-2 (second to last)
• Row n-1: column 0 (first column)

3. Attempting to Optimize with Set or Dictionary

The Pitfall: Some might try to "optimize" by collecting all diagonal positions first using sets to automatically handle duplicates:

# Overcomplicated and less efficient
def diagonalSum_overcomplicated(mat):
    n = len(mat)
    diagonal_positions = set()
  
    # Collect all diagonal positions
    for i in range(n):
        diagonal_positions.add((i, i))
        diagonal_positions.add((i, n-i-1))
  
    # Sum the values
    return sum(mat[r][c] for r, c in diagonal_positions)

While this works, it's actually less efficient due to the overhead of set operations and tuple creation.

The Solution: Stick with the simple linear traversal. The single-pass approach with a conditional check is both clearer and more efficient with O(n) time and O(1) space complexity.