2D Max Pooling That Also Returns the Coordinates of Each Window Maximum

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

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: []

Max Pooling with Argmax Positions

Hints

1. 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.
2. 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?
3. Record the absolute grid coordinates of the best cell, not its offset inside the window.

Approach

The reference first checks whether any window fits: if k is larger than the number of rows or columns it returns []. Otherwise the number of window positions along each axis is (R - k) // stride + 1 and (C - k) // stride + 1, because the top-left corner i * stride must satisfy i * stride + k <= R. For every window (i, j) it starts from the window's top-left cell and scans the k-by-k block in row-major order (increasing row, and increasing column within a row). It replaces the current best only when it sees a strictly larger value, so when the maximum appears several times the first occurrence in row-major order is kept, which is exactly the smallest row and then the smallest column. It stores the absolute grid coordinates (row, col) alongside the value, producing [max_value, row, col] for each output cell. Values are only compared and copied, never added, so 32-bit integers suffice. The direct scan costs O(k^2) per window; with at most 100x100 inputs this is well within limits (the worst case, k around 50 with stride 1, is about 6.5 million comparisons).

• Time complexity: O(out_rows * out_cols * k^2), where out_rows = (R - k) // stride + 1 and out_cols = (C - k) // stride + 1
• Space complexity: O(out_rows * out_cols) for the output; O(1) extra