2D Max Pooling That Also Returns the Coordinates of Each Window Maximum
Company: Waymo
Role: Software Engineer
Category: Coding & Algorithms
Difficulty: medium
Interview Round: Onsite
Implement 2D max pooling over a single-channel grid of integers. A `k`-by-`k` window slides over the grid with step `stride` in both directions, and each window produces the maximum value inside it. As the follow-up, also report where each maximum came from: for every output cell, return the maximum value together with its row and column in the input grid.
### Function Signature
```python
def max_pool_with_argmax(grid: list[list[int]], k: int, stride: int) -> list[list[list[int]]]:
```
### Rules
- No padding is used. A window is placed with its top-left corner at `(i * stride, j * stride)` for every `i, j >= 0` such that the whole window fits inside the grid. With `R` rows and `C` columns, the output has `(R - k) // stride + 1` rows and `(C - k) // stride + 1` columns. Windows that would extend past the grid edge are not produced.
- Output cell `[i][j]` is the list `[max_value, row, col]`, where `max_value` is the largest value in window `(i, j)` and `(row, col)` is its position in `grid`.
- If the maximum occurs more than once in a window, report the occurrence with the smallest row, and among those the smallest column.
- If `k` is larger than the number of rows or columns, no window fits: return an empty list `[]`.
### Constraints
- `1 <= len(grid) <= 100` and `1 <= len(grid[r]) <= 100`; all rows have the same length.
- `-10^9 <= grid[r][c] <= 10^9`
- `1 <= k <= 100`
- `1 <= stride <= 100`
### Examples
**Example 1**
- Input: `grid = [[1, 3, 2, 0], [4, 6, 5, 1], [7, 2, 9, 8], [3, 4, 1, 6]]`, `k = 2`, `stride = 2`
- Output: `[[[6, 1, 1], [5, 1, 2]], [[7, 2, 0], [9, 2, 2]]]`
- Explanation: The four non-overlapping 2-by-2 windows have maxima 6 at `(1, 1)`, 5 at `(1, 2)`, 7 at `(2, 0)` and 9 at `(2, 2)`.
**Example 2**
- Input: `grid = [[5, 1, 5], [2, 5, 0], [1, 1, 3]]`, `k = 2`, `stride = 1`
- Output: `[[[5, 0, 0], [5, 0, 2]], [[5, 1, 1], [5, 1, 1]]]`
- Explanation: The top-left window contains 5 at `(0, 0)` and at `(1, 1)`; the smaller row wins. The top-right window contains 5 at `(0, 2)` and `(1, 1)`; row 0 wins.
**Example 3**
- Input: `grid = [[1, 2], [3, 4]]`, `k = 3`, `stride = 1`
- Output: `[]`
Overview: Implement 2D max pooling with a k-by-k window and a stride, and for each window also return the row and column of its maximum, breaking ties by smallest row then column. Tests window indexing without padding, output sizing, and argmax bookkeeping.
Implement 2D max pooling over a single-channel grid of integers, and report where each maximum came from.
A `k`-by-`k` window slides over `grid` with step `stride` in both directions. Each window produces its maximum value **and** the row and column of that maximum in the input grid.
Implement `max_pool_with_argmax(grid, k, stride)`.
### Rules
- No padding is used. A window is placed with its top-left corner at `(i * stride, j * stride)` for every `i, j >= 0` such that the whole window fits inside the grid. With `R` rows and `C` columns, the output has `(R - k) // stride + 1` rows and `(C - k) // stride + 1` columns. Windows that would extend past the grid edge are not produced.
- Output cell `[i][j]` is the list `[max_value, row, col]`, where `max_value` is the largest value in window `(i, j)` and `(row, col)` is its position in `grid` (0-indexed).
- If the maximum occurs more than once in a window, report the occurrence with the smallest row, and among those the smallest column.
- If `k` is larger than the number of rows or the number of columns, no window fits: return an empty list `[]`.
- The output is compared exactly: output rows in order of `i`, cells within a row in order of `j`.
### Example 1
```
Input: grid = [[1, 3, 2, 0], [4, 6, 5, 1], [7, 2, 9, 8], [3, 4, 1, 6]], k = 2, stride = 2
Output: [[[6, 1, 1], [5, 1, 2]], [[7, 2, 0], [9, 2, 2]]]
```
The four non-overlapping 2-by-2 windows have maxima 6 at `(1, 1)`, 5 at `(1, 2)`, 7 at `(2, 0)` and 9 at `(2, 2)`.
### Example 2
```
Input: grid = [[5, 1, 5], [2, 5, 0], [1, 1, 3]], k = 2, stride = 1
Output: [[[5, 0, 0], [5, 0, 2]], [[5, 1, 1], [5, 1, 1]]]
```
The top-left window contains 5 at `(0, 0)` and at `(1, 1)`; the smaller row wins. The top-right window contains 5 at `(0, 2)` and `(1, 1)`; row 0 wins.
### Example 3
```
Input: grid = [[1, 2], [3, 4]], k = 3, stride = 1
Output: []
```
### Constraints
- `1 <= len(grid) <= 100` and `1 <= len(grid[r]) <= 100`; all rows have the same length.
- `-10^9 <= grid[r][c] <= 10^9` (every value fits in a signed 32-bit integer).
- `1 <= k <= 100`
- `1 <= stride <= 100`
Constraints
- 1 <= len(grid) <= 100 and 1 <= len(grid[r]) <= 100; all rows have the same length.
- -10^9 <= grid[r][c] <= 10^9 (every value fits in a signed 32-bit integer).
- 1 <= k <= 100
- 1 <= stride <= 100
Examples
Input: ([[1, 3, 2, 0], [4, 6, 5, 1], [7, 2, 9, 8], [3, 4, 1, 6]], 2, 2)
Expected Output: [[[6, 1, 1], [5, 1, 2]], [[7, 2, 0], [9, 2, 2]]]
Explanation: Example 1: four non-overlapping 2x2 windows.
Input: ([[5, 1, 5], [2, 5, 0], [1, 1, 3]], 2, 1)
Expected Output: [[[5, 0, 0], [5, 0, 2]], [[5, 1, 1], [5, 1, 1]]]
Explanation: Example 2: repeated maximum resolved by smallest row, then smallest column.
Hints
- Work out the output shape first: how many top-left corners i * stride keep a k-row window inside R rows? If k exceeds either dimension there are none.
- For each window, scan its cells in row-major order (row by row, left to right). Which comparison, strict or non-strict, keeps the first occurrence of a repeated maximum?
- Record the absolute grid coordinates of the best cell, not its offset inside the window.