Corner to Corner

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

Domain
Python
Difficulty
medium
Seniority
L4

Problem

You are given cells, a square matrix of integers from a tabular export flattened row by row into one list, and size, the matrix's side length. Return the values on the main diagonal, from the top-left corner to the bottom-right.

A grid pressed flat still remembers its corners.

Example 1

Input:

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

Output:

[1,5,9]

Example 2

Input:

cells = [10,20,30,40], size = 2

Output:

[10,40]

Worked solution and explanation

What this is really testing

Under the matrix costume this is stride arithmetic, not a grid walk. The grid has already been pressed into one flat line, so there is no row-of-cells to iterate and no i == j to check: the values you want are simply the ones whose flat index is i*size + i. Candidates who see that write one line. Candidates who don't reach for cells[i] and hand back the first row, or scan the whole buffer hunting for diagonal values by content, which cannot work because a number carries no memory of where it sat. Get the stride wrong and a 1x1 buffer still passes by luck, so the mistake hides until a bigger matrix exposes it.

Where the diagonal went

Step 1: From (i, j) to a single offset

Row-major flattening lays row 0 first, then row 1, and so on. So the element at grid position (i, j) lands at flat index i*size + j. The main diagonal is every cell where the column equals the row, j equals i, which pins its flat index at i*size + i. That is the whole trick: one expression turns a 2-D coordinate into a 1-D offset.

Step 2: One counter, computed offset

One counter drives it. As i runs from 0 to size-1 you read cells[i*size + i] each time and never touch an off-diagonal element. Notice what is NOT available: iterating cells directly gives you a stream of integers with no way to tell a diagonal value from any other, so the side length is what makes the selection possible at all.

The solution

Stride-indexed extraction of the main diagonal
def extract_diagonal(cells, size):
    return [cells[i * size + i] for i in range(size)]

One offset per row, computed straight from the side length.

Common follow-up questions