Find Median from Data Stream

▶

Steps

setup

Initialize MedianFinder

Create empty max heap 'small' and min heap 'large', initialize delayed removals map and size counters.

💡 Starting with empty heaps and counters sets the foundation for adding numbers and balancing heaps.

Line:

self.small = []  # max heap (invert values)
self.large = []  # min heap
self.delayed = {}
self.small_size = 0
self.large_size = 0

💡 Heaps and counters are ready to track numbers and maintain balance.

insert

Add number 1 to 'small' heap

Since 'small' is empty, add 1 to 'small' as -1 (max heap stores negatives). Increment small_size.

💡 First number goes to 'small' heap to start building the data structure.

Line:

if not self.small or num <= -self.small[0]:
    heapq.heappush(self.small, -num)
    self.small_size += 1

💡 The first number always goes to 'small' heap to start balancing.

insert

Add number 2 to 'large' heap

Number 2 is greater than -small[0] (which is 1), so add 2 to 'large' heap and increment large_size.

💡 Numbers greater than the max of 'small' go to 'large' heap to keep order.

Line:

else:
    heapq.heappush(self.large, num)
    self.large_size += 1

💡 Splitting numbers between heaps keeps them balanced and ordered.

shrink

Balance heaps after adding 2

Heaps are balanced: small_size=1, large_size=1, difference is 0, so no rebalancing needed.

💡 Balancing ensures heaps differ in size by at most 1 for correct median calculation.

Line:

if self.small_size > self.large_size + 1:
    ...
elif self.large_size > self.small_size:
    ...

💡 Balanced heaps mean median can be computed from tops directly.

compare

Find median after adding 1 and 2

Since both heaps have equal size, median is average of tops: (1 + 2) / 2 = 1.5.

💡 Median for even count is average of max of 'small' and min of 'large'.

Line:

if self.small_size == self.large_size:
    median = (-self.small[0] + self.large[0]) / 2

💡 Median depends on top elements of both heaps when sizes match.

insert

Add number 3 to 'large' heap

Number 3 is greater than -small[0] (1), so add 3 to 'large' heap and increment large_size.

💡 New number goes to 'large' because it is greater than max of 'small'.

Line:

else:
    heapq.heappush(self.large, num)
    self.large_size += 1

💡 Adding to 'large' may unbalance heaps, requiring rebalancing.

shrink

Balance heaps after adding 3

Since large_size (2) > small_size (1), move smallest from 'large' (2) to 'small' as -2. Update sizes accordingly.

💡 Balancing keeps heaps sizes within 1 to maintain median correctness.

Line:

elif self.large_size > self.small_size:
    val = heapq.heappop(self.large)
    self.large_size -= 1
    heapq.heappush(self.small, -val)
    self.small_size += 1

💡 Rebalancing moves elements between heaps to maintain size property.

compare

Find median after adding 3

Now small_size (2) > large_size (1), median is top of 'small' heap: 2 (negate -2).

💡 Median for odd count is top of larger heap.

Line:

else:
    median = -self.small[0]

💡 Median depends on heap sizes and top elements.

reconstruct

Summary: Median after first findMedian call

Median after adding 1 and 2 is 1.5, computed from tops of both heaps.

💡 Shows how median is derived from balanced heaps with even count.

Line:

return median

💡 Median calculation depends on heap sizes and top elements.

reconstruct

Summary: Median after second findMedian call

Median after adding 3 is 2, the top of 'small' heap which has more elements.

💡 Shows median for odd count from the larger heap top.

Line:

return median

💡 Median is the middle element when count is odd.

import heapq

class MedianFinder:
    def __init__(self):
        self.small = []  # max heap (invert values)  # STEP 1
        self.large = []  # min heap  # STEP 1
        self.delayed = {}  # STEP 1
        self.small_size = 0  # STEP 1
        self.large_size = 0  # STEP 1

    def _prune(self, heap):
        while heap and self.delayed.get(abs(heap[0]), 0):
            num = abs(heap[0])
            self.delayed[num] -= 1
            if self.delayed[num] == 0:
                del self.delayed[num]
            heapq.heappop(heap)

    def addNum(self, num: int) -> None:
        if not self.small or num <= -self.small[0]:  # STEP 2,6
            heapq.heappush(self.small, -num)  # STEP 2,6
            self.small_size += 1  # STEP 2,6
        else:
            heapq.heappush(self.large, num)  # STEP 3,6
            self.large_size += 1  # STEP 3,6

        # Balance heaps
        if self.small_size > self.large_size + 1:  # STEP 7
            val = -heapq.heappop(self.small)  # STEP 7
            self.small_size -= 1  # STEP 7
            heapq.heappush(self.large, val)  # STEP 7
            self.large_size += 1  # STEP 7
        elif self.large_size > self.small_size:  # STEP 7
            val = heapq.heappop(self.large)  # STEP 7
            self.large_size -= 1  # STEP 7
            heapq.heappush(self.small, -val)  # STEP 7
            self.small_size += 1  # STEP 7

    def findMedian(self) -> float:
        if self.small_size == self.large_size:  # STEP 5
            return (-self.small[0] + self.large[0]) / 2  # STEP 5
        else:
            return -self.small[0]  # STEP 8

    def removeNum(self, num: int) -> None:
        self.delayed[num] = self.delayed.get(num, 0) + 1
        if self.small and num <= -self.small[0]:
            self.small_size -= 1
            if num == -self.small[0]:
                self._prune(self.small)
        else:
            self.large_size -= 1

📊

Find Median from Data Stream - Watch the Algorithm Execute, Step by Step

Watching the algorithm step-by-step reveals how the two heaps maintain a balanced state to efficiently find the median at any point in the data stream.

setup

insert

small_top

-1

0

insert

small_top

-1

0

large_top

2

1

compare

small_top

-1

0

large_top

2

1

compare

small_top

-1

0

large_top

2

1

Result: 1.5

insert

small_top

-1

0

large_top

2

1

3

2

Result: 1.5

rebalance

small_top

-2

0

-1

1

large_top

3

2

Result: 1.5

compare

small_top

-2

0

-1

1

large_top

3

2

Result: 2

record

small_top

-1

0

large_top

2

1

Result: 1.5

record

small_top

-2

0

-1

1

large_top

3

2

Result: 2

Key Takeaways

✓ Two heaps maintain a balanced partition of the data stream to efficiently compute median.

This balance is hard to see from code alone but is clear when watching heaps grow and shrink.

✓ Median calculation depends on heap sizes: average of tops if equal, top of larger heap if odd count.

Visualizing heap sizes and top elements clarifies why median formulas differ.

✓ Rebalancing moves elements between heaps to maintain size constraints after each insertion.

Seeing the actual element moved during rebalance helps understand the heap size invariants.

Practice

1. You need to design a data structure that supports adding intervals, removing intervals, and querying whether a given range is fully covered by the added intervals. Which approach guarantees efficient updates and queries with logarithmic time complexity per operation?

easy

A. Use a balanced binary search tree (e.g., TreeMap) to store and merge intervals, enabling logarithmic time operations.

B. Use a brute force list of intervals and scan all intervals for each operation.

C. Use dynamic programming to precompute coverage for all possible ranges.

D. Use a greedy algorithm that always merges intervals only when queried.

Solution

1. Step 1: Understand the problem requirements

   The data structure must support add, remove, and query operations on intervals efficiently.

2. Step 2: Evaluate approaches

   Brute force scanning is O(n) per operation, which is inefficient. DP is not suitable for dynamic interval updates. Greedy merging only on query delays merging and leads to inefficient queries. Balanced BST with interval merging maintains sorted intervals and merges overlapping ones, supporting O(log n) operations.

3. Final Answer:

   Option A -> Option A

4. Quick Check:

   Balanced BST with merging -> O(log n) per operation [OK]

Hint: Balanced BST with merging enables efficient interval operations [OK]

Common Mistakes:

• Thinking DP or greedy solves dynamic interval updates efficiently

2. Consider this buggy snippet of the optimal Twitter design's getNewsFeed method:

def getNewsFeed(self, userId: int) -> list:
    users = self.followees.get(userId, set()).copy()
    heap = []  # max heap by time

    for u in users:
        head = self.userTweets.get(u, None)
        if head:
            heapq.heappush(heap, (-head.time, head))

    result = []
    while heap and len(result) < 10:
        time, node = heapq.heappop(heap)
        result.append(node.tweetId)
        if node.next:
            heapq.heappush(heap, (-node.next.time, node.next))
    return result

What is the subtle bug in this code?

medium

A. The user's own tweets are not included because userId is not added to users set.

B. The heap is used as a min heap instead of max heap by pushing negative times.

C. The linked list nodes are not updated correctly after popping from the heap.

D. The followees set is copied unnecessarily causing performance overhead.

Solution

1. Step 1: Check inclusion of user's own tweets

   The code gets followees but never adds the userId itself to the users set, so user's own tweets are excluded.

2. Step 2: Verify heap usage and linked list traversal

   Negative times are correctly used for max heap behavior; linked list next pointers are used properly.

3. Final Answer:

   Option A -> Option A

4. Quick Check:

   Missing userId in users set causes incorrect feed [OK]

Hint: User's own tweets must be explicitly included in feed [OK]

Common Mistakes:

• Forgetting to add userId to followees set
• Misusing heap sign for max heap

3. What is the amortized time complexity of the put operation in the optimal LFU Cache implementation using a frequency map of doubly linked lists and hash maps?

medium

A. O(n), where n is the number of keys in the cache

B. O(f), where f is the number of distinct frequencies

C. O(log n), due to maintaining frequency order

D. O(1), amortized constant time

Solution

1. Step 1: Analyze data structures used

   Hash maps provide O(1) access to keys and frequencies. Doubly linked lists allow O(1) insertion and deletion of keys within frequency groups.

2. Step 2: Consider operations in put

   Eviction and frequency updates involve O(1) operations on hash maps and linked lists. Frequency increments and min frequency updates are also O(1).

3. Final Answer:

   Option D -> Option D

4. Quick Check:

   Optimal LFU cache achieves O(1) amortized put [OK]

Hint: Hash maps + linked lists enable O(1) put [OK]

Common Mistakes:

• Assuming O(n) due to scanning keys
• Thinking frequency updates cost O(log n)
• Confusing distinct frequencies with n

4. What is the time complexity of the divide and conquer approach to the Skyline Problem when given n buildings? Assume merging two skylines of total length m takes O(m) time.

medium

A. O(n²)

B. O(n)

C. O(n log n)

D. O(n * w) where w is the width of the skyline

Solution

1. Step 1: Identify divide and conquer recurrence

   The algorithm splits buildings into halves recursively, then merges skylines in O(m) time where m is proportional to n.

2. Step 2: Solve recurrence and analyze merge cost

   Recurrence is T(n) = 2T(n/2) + O(n), which solves to O(n log n). The width w does not affect complexity here.

3. Final Answer:

   Option C -> Option C

4. Quick Check:

   Divide and conquer with linear merge per level -> O(n log n) [OK]

Hint: Divide and conquer recurrence T(n)=2T(n/2)+O(n) -> O(n log n) [OK]

Common Mistakes:

• Confusing width w with n
• Assuming quadratic due to nested loops

5. Suppose you want to extend the LRU Cache to support a "reuse" operation that marks an existing key as recently used without retrieving its value. Which modification to the optimal LRU Cache implementation is correct to support this efficiently?

hard

A. On reuse, remove the key from the cache and re-insert it with the same value.

B. Do nothing; reuse is equivalent to get, so no code change is needed.

C. Use a queue instead of a doubly linked list to track usage order for easier reuse.

D. Add a new method that calls _remove and _add on the node corresponding to the key.

Solution

1. Step 1: Understand reuse operation

   Reuse means marking a key as recently used without returning its value, so usage order must be updated.

2. Step 2: Efficiently update usage order

   Calling _remove and _add on the node moves it to the front in O(1) time, correctly updating usage.

3. Final Answer:

   Option D -> Option D

4. Quick Check:

   Reuse requires usage update without value retrieval, done by repositioning node [OK]

Hint: Reuse updates usage order like get but without returning value [OK]

Common Mistakes:

• Re-inserting wastes time
• Queue can't reorder efficiently
• Assuming reuse needs no code change