Katkat
Back

Compute Max Pooling Values and Coordinates

WaymoSoftware Engineer · Technical Screen
MediumSoftware EngineerTechnical ScreenCoding & Algorithms

Last updated

This problem tests your ability to manipulate 2D arrays, manage sliding window bounds, and track metadata (indices) efficiently during a data transformation task.

Max Pooling with Coordinates

You are tasked with implementing a max_pool function that performs a 2D max pooling operation on a grid. In addition to the maximum value in each window, you must return the original (row, col) coordinates of that maximum value.

Function Signature

def max_pool(grid, kernel_size, stride):
    # grid: List[List[int]]
    # kernel_size: int
    # stride: int
    # returns: List[List[Tuple[int, Tuple[int, int]]]]

Specifications

  • The function takes a 2D integer grid, the size of the square pooling window (kernel_size), and the stride length.
  • For each window, find the maximum value. If there is a tie, choose the coordinate that appears first in row-major order (top-to-bottom, then left-to-right).
  • The output should be a 2D list of tuples, where each element is (max_value, (row, col)).
  • Assume the input grid is always rectangular and the kernel fits within the bounds given the stride.

Example

Input: grid = [[1, 3, 2], [4, 5, 6], [7, 8, 9]], kernel_size = 2, stride = 1

Output: [[[5, (1, 1)], [6, (1, 2)]], [[8, (2, 1)], [9, (2, 2)]]]

Explanation:

  • Window 1 (top-left 2x2): [[1, 3], [4, 5]]. Max is 5 at (1, 1).
  • Window 2 (top-right 2x2): [[3, 2], [5, 6]]. Max is 6 at (1, 2).
  • Window 3 (bottom-left 2x2): [[4, 5], [7, 8]]. Max is 8 at (2, 1).
  • Window 4 (bottom-right 2x2): [[5, 6], [8, 9]]. Max is 9 at (2, 2).

Edge Cases

  • What if the kernel size is 1?
  • What if the stride is equal to the kernel size?
  • How should the function behave if the grid dimensions are not perfectly divisible by the stride?

Hints

Hint 1

Calculate the number of output rows and columns first using the formula: output_dim = (input_dim - kernel_size) // stride + 1.

Hint 2

Use nested loops: the outer two loops iterate over the top-left corner coordinates of each window, while the inner two loops iterate through the elements within that specific window.

Hint 3

Maintain a variable to track the current maximum value and its coordinates during the iteration of each window to avoid redundant passes.

Follow-up questions

  • How would you modify the function if the stride were different for rows and columns?
  • What is the time complexity of your approach in terms of the grid size (N x M) and kernel size (K x K)? Can it be optimized using a monotonic queue?
  • If the grid were extremely large and didn't fit in memory, how would you process it?