The Diagonal Accountant

An easy Python interview practice problem on DataDriven. Write the Python and run it against test cases, with instant feedback.

Domain
Python
Difficulty
easy
Seniority
L3
Asked in
Citadel

Problem

Given an n x n matrix, return the sum of main-diagonal and anti-diagonal values. For odd n, do NOT double-count the center element.

Two diagonals cross in the center of every square.

Example 1

Input:

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

Output:

25

Example 2

Input:

matrix = [[1,1,1,1],[1,1,1,1],[1,1,1,1],[1,1,1,1]]

Output:

8

Worked solution and explanation

Why this problem exists in real interviews

This tests matrix traversal with index arithmetic, specifically accessing diagonal elements. It probes whether a candidate can derive the index patterns for main and anti-diagonals and handle the overlap at the center of odd-sized matrices.

Break down the requirements

Step 1: Sum the main diagonal

Elements at positions (i, i) for i from 0 to n-1.

Step 2: Sum the anti-diagonal

Elements at positions (i, n-1-i) for i from 0 to n-1.

Step 3: Subtract the center if n is odd

When n is odd, position (n//2, n//2) is on both diagonals and gets counted twice.

The solution

Index-based diagonal traversal
def diagonal_sum(matrix: list) -> int:
    n = len(matrix)
    total = 0
    for i in range(n):
        total += matrix[i][i]
        total += matrix[i][n - 1 - i]
    if n % 2 == 1:
        total -= matrix[n // 2][n // 2]
    return total

Common follow-up questions