Join Our 5-Week ML/AI Engineer Interview Bootcamp 🚀 led by ML Tech Leads at FAANGs

Back to Questions

67. Pooling operations

medium
GeneralGeneral
staff

Implement 2D pooling (max and average) over a single-channel image, a common downsampling step in computer vision models. You’ll take an input image as a NumPy array of shape (H,W)(H, W) and return the pooled output as a NumPy array.

Pooling is defined over each window as:

yi,j=max⁡0≤a<kh, 0≤b<kwxi⋅sh+a, j⋅sw+boryi,j=1khkw∑a=0kh−1∑b=0kw−1xi⋅sh+a, j⋅sw+by_{i,j} = \max_{0 \le a < k_h,\ 0 \le b < k_w} x_{i\cdot s_h + a,\ j\cdot s_w + b} \quad\text{or}\quad y_{i,j} = \frac{1}{k_h k_w}\sum_{a=0}^{k_h-1}\sum_{b=0}^{k_w-1} x_{i\cdot s_h + a,\ j\cdot s_w + b}

Requirements

Implement the function

python

Rules:

  • Use only NumPy and Python built-in libraries (no SciPy, OpenCV, etc.).
  • mode="max" must compute max-pooling; mode="avg" must compute average-pooling.
  • Use valid pooling only (no padding): only windows fully inside the image are included.
  • Return the result as a NumPy array.
  • Keep it in a single function.

Example

python

Output:

python
Input Signature
ArgumentType
modestr
imagenp.ndarray
stridelist
kernel_sizelist
Output Signature
Return NameType
valuenp.ndarray

Constraints

  • Valid pooling only; no padding.

  • Return NumPy array.

  • Use NumPy + built-ins only.

Hint 1

Compute output shape for valid pooling: out_h = (H-k_h)//s_h + 1, out_w = (W-k_w)//s_w + 1.

Hint 2

For each output (i,j), slice the window with top-left (i*s_h, j*s_w) and shape (k_h,k_w); apply max or mean.

Hint 3

Input image is np.ndarray. Initialize output with np.zeros. Check mode is {max, avg}.

Roles
ML Engineer
AI Engineer
Companies
GeneralGeneral
Levels
staff
senior
entry
Tags
numpy-slicing
2d-pooling
sliding-window
computer-vision
48 people are solving this problem
Python LogoPython Editor
Ln 1, Col 1

Input Arguments

Edit values below to test with custom inputs

You need tolog in/sign upto run or submit