Solve sliding-window and heap problems
Company: Amazon
Role: Software Engineer
Category: Coding & Algorithms
Difficulty: Medium
Interview Round: Onsite
##### Question
Design an algorithm using a sliding window to compute the maximum sum of any contiguous subarray of length k; explain time- and space-complexity optimizations and how you would adapt the solution for concurrent execution. Maintain the median of the last k numbers of a real-time integer stream; choose appropriate data structures (e.g., two heaps, hashmap) and discuss trade-offs and optimizations.
Quick Answer: This question evaluates algorithm design and data structure proficiency, focusing on sliding-window techniques for fixed-length subarray aggregates and heap- or hashmap-based approaches for maintaining medians in a real-time integer stream, along with considerations for concurrent execution.
Given an integer array nums and an integer k, return a list of the medians of every contiguous subarray (sliding window) of length k. The median of a window is the middle value after sorting the k elements; if k is even, it is the average of the two middle values. Return each median as a float, in the order of the windows from left to right.
Constraints
- 1 <= len(nums) <= 200000
- 1 <= k <= len(nums)
- -10^9 <= nums[i] <= 10^9
- Return medians as floats; for even k, use (a + b) / 2.0 with no rounding beyond floating-point defaults
- Target time complexity: O(n log k), space: O(k)
Solution
from typing import List
import heapq
from collections import defaultdict
def sliding_window_median(nums: List[int], k: int) -> List[float]:
small = [] # max-heap (store negatives)
large = [] # min-heap
delayed = defaultdict(int) # value -> count to lazily delete
s_size = 0 # number of valid elements in small
l_size = 0 # number of valid elements in large
def prune(heap):
# Remove elements at the top that are marked for deletion
while heap:
if heap is small:
num = -heap[0]
else:
num = heap[0]
if delayed.get(num, 0):
heapq.heappop(heap)
delayed[num] -= 1
if delayed[num] == 0:
del delayed[num]
else:
break
def make_balance():
nonlocal s_size, l_size
prune(small)
prune(large)
# Ensure size property: s_size >= l_size and at most 1 larger
if s_size > l_size + 1:
val = -heapq.heappop(small)
heapq.heappush(large, val)
s_size -= 1
l_size += 1
elif s_size < l_size:
val = heapq.heappop(large)
heapq.heappush(small, -val)
s_size += 1
l_size -= 1
prune(small)
prune(large)
def add_num(num: int):
nonlocal s_size, l_size
if not small or num <= -small[0]:
heapq.heappush(small, -num)
s_size += 1
else:
heapq.heappush(large, num)
l_size += 1
make_balance()
def remove_num(num: int):
nonlocal s_size, l_size
delayed[num] += 1
# Decide which heap the number belongs to
if small and num <= -small[0]:
s_size -= 1
if small and -small[0] == num:
prune(small)
else:
l_size -= 1
if large and large[0] == num:
prune(large)
make_balance()
def get_median() -> float:
prune(small)
prune(large)
if k % 2 == 1:
return float(-small[0])
else:
return (-small[0] + large[0]) / 2.0
res: List[float] = []
for i, num in enumerate(nums):
add_num(num)
if i >= k:
remove_num(nums[i - k])
if i >= k - 1:
res.append(get_median())
return res
Explanation
Maintain two heaps representing the lower and upper halves of the current window: a max-heap (as negatives) for the lower half and a min-heap for the upper half. Insert each new element into the appropriate heap, and lazily delete elements that slide out using a hashmap (value -> count). Balance heaps so that the max-heap has either the same count as the min-heap or one more valid element. The median is either the top of the max-heap (odd k) or the average of both heap tops (even k). Pruning removes any delayed elements from heap tops just before they are used or moved. This approach yields O(n log k) time with O(k) space.
Time complexity: O(n log k). Space complexity: O(k).
Hints
- Maintain two heaps: a max-heap for the lower half and a min-heap for the upper half of the window.
- Balance heaps so that the max-heap has either the same number of valid elements as the min-heap or one more.
- Use a hashmap for lazy deletion to remove elements that slide out without direct heap deletion.
- When computing the median, prune invalid (delayed) elements from heap tops.
- For even k, average the max of the lower half and the min of the upper half.