# Corner to Corner > A grid pressed flat still remembers its corners. Canonical URL: Domain: Python · Difficulty: medium · Seniority: L4 ## Problem A square matrix from a tabular export arrives flattened: a single row-major list `cells` holding n*n integers, plus its side length `size`. Return the values on the main diagonal, from the top-left corner to the bottom-right. An empty buffer returns an empty list. ## 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** ```python 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.* > **Why the direct index wins** > > Time is O(size): you compute one index and read one element per diagonal position, never the off-diagonal cells. Space is O(size) for the output. Reconstructing the full 2-D matrix first, or looping over every flat position and testing it, is O(size*size) for the same answer, and an interviewer notices the wasted work on a large buffer. > **Common pitfall** > > Reading cells[i] instead of cells[i*size + i]. It looks plausible and it even passes a 1x1 buffer, where both expressions equal cells[0], so the bug survives the smallest test. But cells[i] walks the first row, not the diagonal, and breaks on any size above one. The other trap is scanning the buffer by value: a flattened integer no longer knows its coordinates, so content-based selection is impossible. > **In production** > > Flat buffers plus a shape are how numeric arrays actually travel between systems. A raw export, a serialized tensor, or a memory-mapped column store hands you bytes in row-major order and a shape tuple; every access into it is offset arithmetic against the stride. Getting comfortable with index = row*width + col is the same muscle you use when you stop materializing intermediate matrices and read straight from the buffer. --- ## Common follow-up questions - How would you pull the anti-diagonal, top-right to bottom-left, from the same flat buffer? _(Tests flipping the offset to i*size + (size - 1 - i).)_ - What changes if the buffer is a rows-by-cols rectangle instead of a square? _(Tests generalizing the stride to a non-square shape, where the diagonal length is min(rows, cols) and the offset uses the column count as the width.)_ - How would you return every diagonal running in the same direction, not just the main one? _(Tests iterating over diagonal offsets: cells on one diagonal share a constant value of (col - row), so each offset shifts the starting column.)_ ## Related - [All practice problems](https://datadriven.io/problems) - [Mock interview mode](https://datadriven.io/interview/corner_to_corner) - [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). DataDriven is the data engineering interview community. Live code execution in SQL, Python, and Spark sandboxes. Every feature is open to every member.