# Diagonal Extract

> Not every value sits in a row or column.

Canonical URL: <https://datadriven.io/problems/diagonal_extract>

Domain: Python · Difficulty: medium · Seniority: L4

## Problem

Given a square 2D matrix, return a list of the main-diagonal values (matrix[i][i] for i in range(n)).

## Worked solution and explanation

### Why this problem exists in real interviews

Extracting the main diagonal of a matrix tests **index-based access** and understanding of 2D list structure. It is a quick check on whether you can map the concept of `row == column` to a loop.

---

### Break down the requirements

#### Step 1: Identify diagonal elements

The main diagonal consists of elements where the row index equals the column index: `matrix[i][i]`.

#### Step 2: Handle non-square matrices

The diagonal length is `min(rows, cols)`. Stop at whichever dimension is shorter.

---

### The solution

**Index-aligned extraction of the main diagonal**

```python
def extract_diagonal(matrix):
    result = []
    n = len(matrix)
    for i in range(n):
        if i < len(matrix[i]):
            result.append(matrix[i][i])
    return result
```

> **Time and Space Complexity**
>
> **Time:** O(min(r, c)) where r is the number of rows and c is the number of columns. Only diagonal elements are visited.
> 
> **Space:** O(min(r, c)) for the output list.

> **Interviewers Watch For**
>
> Using a single loop with `matrix[i][i]` rather than nested loops. The diagonal is fully determined by one index, so nested iteration is unnecessary overhead.

> **Common Pitfall**
>
> Assuming the matrix is always square. Non-square matrices (more rows than columns or vice versa) require a bounds check on both dimensions.

---

## Common follow-up questions

- How would you extract the anti-diagonal? _(Tests using `matrix[i][n - 1 - i]` where n is the column count.)_
- What if the matrix is stored as a flat list? _(Tests index arithmetic: the diagonal element for row i in a width-w matrix is at position `i * w + i`.)_
- How would you sum all diagonals of a matrix? _(Tests iterating over all possible diagonal offsets, not just the main one.)_

## Related

- [All practice problems](https://datadriven.io/problems)
- [Mock interview mode](https://datadriven.io/interview/diagonal_extract)
- [Python Interview Questions](https://datadriven.io/python-interview-questions)
- [Data Engineering Interview Prep Guide](https://datadriven.io/data-engineer-interview-prep)
- [Daily Challenge](https://datadriven.io/daily)

---

Source: DataDriven (https://datadriven.io). 100% free data engineering interview prep. Live code execution against Postgres 16, Python 3.11, and Spark sandboxes. No paywall, no premium tier, no signup gate.