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