Heap Hard Problems: Kth Element in Matrix (5 Approaches) & Find Median in a Stream

This chapter explores two classic, high-frequency heap and priority queue problems: Kth Smallest Element in a Row-Column Sorted Matrix and Find Median from a Data Stream. Both problems are broken down from naive brute-force solutions to highly optimized algorithmic strategies, complete with architectural comparisons and production-ready C++ code.

1. Problem 1: Kth Smallest Element in a Matrix (Overview & Status Quo)

Problem Statement

Given an $n \times n$ matrix where each of the rows and columns is sorted in non-decreasing order, find the $k^{\text{th}}$ smallest element in the matrix.

Visualizing the Matrix Structure

Consider a $4 \times 4$ matrix sorted along both rows and columns:

Row 0: 16  22  27  38
Row 1: 28  36  41  50
Row 2: 60  64  78  87
Row 3: 81  89  90  94

• Goal: Find the $6^{\text{th}}$ smallest element.
• Manual Walkthrough:
  • $1^{\text{st}}$ smallest = 16
  • $2^{\text{nd}}$ smallest = 22
  • $3^{\text{rd}}$ smallest = 27
  • $4^{\text{th}}$ smallest = 28
  • $5^{\text{th}}$ smallest = 36
  • $6^{\text{th}}$ smallest = 41 (Target Answer)

2. Approach 1: The Naive Flatten & Sort Strategy (Brute Force)

The most intuitive brute-force strategy is to copy all elements of the $n \times n$ matrix into a $1\text{D}$ array, sort the array, and pick the element at index $k - 1$.

Algorithm Steps

1. Create a $1\text{D}$ array or vector of size $n^2$.
2. Copy all matrix elements into this $1\text{D}$ container.
3. Sort the container using Merge Sort, Quick Sort, or standard library sorting.
4. Return the element at index k - 1.

Complexity Analysis

• Time Complexity: $\mathcal{O}(n^2 \log n^2) = \mathcal{O}(n^2 \log n)$ due to sorting $n^2$ elements.
• Space Complexity: $\mathcal{O}(n^2)$ to store all elements in the $1\text{D}$ array.

3. Approach 2: Naive Min-Heap Strategy

Instead of fully sorting the array, we can insert all elements into a Min-Heap (Priority Queue) and extract the minimum element $k$ times.

Architectural Misconception Warning

Note: Many tutorials mistakenly illustrate priority queues using a sorted vertical stack. A priority queue backed by a binary heap does not store elements in a strictly sorted linear stack; elements are organized as a complete binary tree where parent nodes are smaller than their children. Avoid using misleading stack diagrams in interviews.

Algorithm Steps

1. Copy all $n^2$ matrix elements into a $1\text{D}$ array/vector.
2. Build a Min-Heap from the array in $\mathcal{O}(n^2)$ time using the bottom-up heapification method.
3. Perform pop() operations $k$ times. The $k^{\text{th}}$ pop yields the target element.

Complexity Analysis

• Time Complexity: $\mathcal{O}(n^2 + k \log n^2) = \mathcal{O}(n^2 + k \log n)$
• Space Complexity: $\mathcal{O}(n^2)$

4. Approach 3: Optimized Heap Strategy (Leveraging Matrix Sorted Properties)

Approaches 1 and 2 ignore the fact that the matrix is already row-wise and column-wise sorted. We can optimize this by maintaining a Min-Heap of size $n$ containing elements from the first column, then stepping horizontally across rows.

Core Mental Model & Real-Life Analogy

Real-Life Analogy: Imagine parking cars in lanes where each lane is sorted by height. To find the $k^{\text{th}}$ shortest car, you don't line up every single car in a massive lot. Instead, you look at the front car of each lane. Take the shortest car out, and immediately replace it with the next car waiting in that exact same lane.

Step-by-Step Flow Walkthrough

1. Push the first element of each row (along with its row and column indices) into a Min-Heap. Because elements are stored as std::pair<int, std::pair<int, int>>, the heap automatically sorts primarily by the matrix value.
2. Extract the minimum element from the heap (this represents the next smallest element in the matrix).
3. If the extracted element has a valid right neighbor in the same row (col + 1 < n), push that neighbor into the Min-Heap.
4. Repeat this extraction and insertion process $k$ times.

C++ Implementation

#include <vector>
#include <queue>

using namespace std;

class Solution {
public:
    int kthSmallest(vector<vector<int>>& matrix, int k) {
        int n = matrix.size();
        
        // Min-heap storing: {value, {row, col}}
        priority_queue<
            pair<int, pair<int, int>>, 
            vector<pair<int, pair<int, int>>>, 
            greater<pair<int, pair<int, int>>>
        > minHeap;
        
        // Initialize heap with the first element of each row
        for (int i = 0; i < n; i++) {
            minHeap.push({matrix[i][0], {i, 0}});
        }
        
        int ans = 0;
        for (int count = 0; count < k; count++) {
            auto top = minHeap.top();
            minHeap.pop();
            
            ans = top.first;
            int r = top.second.first;
            int c = top.second.second;
            
            // Push the next element in the same row if it exists
            if (c + 1 < n) {
                minHeap.push({matrix[r][c + 1], {r, c + 1}});
            }
        }
        
        return ans;
    }
};

Complexity Analysis

• Time Complexity: $\mathcal{O}(n + k \log n)$ — initializing the heap takes $\mathcal{O}(n)$, and each of the $k$ extractions/insertions takes $\mathcal{O}(\log n)$.
• Space Complexity: $\mathcal{O}(n)$ to maintain at most $n$ elements in the heap.

5. Approaches 4 & 5: Advanced Optimization (Directional Tracking & Binary Search on Range)

Approach 4: Matrix Search via Single Directional Spreading

Instead of inserting the entire first column initially, we can start with the top-left element matrix[0][0] and dynamically expand right and down while maintaining visited tracking to avoid duplicate processing.

Approach 5: Binary Search on Value Range (Parametric Search)

Since the matrix elements are bounded between matrix[0][0] (minimum possible value) and matrix[n-1][n-1] (maximum possible value), we can apply Binary Search on the Answer Range:

1. Set low = matrix[0][0] and high = matrix[n-1][n-1].
2. Compute mid = low + (high - low) / 2.
3. Count how many numbers in the matrix are less than or equal to mid in $\mathcal{O}(n)$ time by traversing from the bottom-left or top-right corner.
4. If the count is less than $k$, our target must be larger, so set low = mid + 1. Otherwise, set high = mid - 1 and record mid as a potential answer.

6. Problem 2: Find Median in a Data Stream (Problem Statement)

Problem Statement

Design a data structure that supports receiving a stream of integers and returning the running median of all elements seen so far.

• Odd Number of Elements: The median is the exact middle element after sorting.
• Even Number of Elements: The median is the average of the two middle elements.

Naive Insertion Approach & Bottleneck

If we use a standard array and insert incoming stream elements into their correct sorted positions, shifting elements takes $\mathcal{O}(n)$ per insertion, leading to an $\mathcal{O}(n^2)$ overall time complexity.

7. The Dual-Heap Architecture (Optimal Solution)

To achieve $\mathcal{O}(\log n)$ insertion and $\mathcal{O}(1)$ median retrieval, we split the incoming data stream into two halves using a Max-Heap and a Min-Heap.

Invariant Rules

1. Left Max-Heap stores the smaller half of the numbers.
2. Right Min-Heap stores the larger half of the numbers.
3. Size Property: The size of the Left Max-Heap must either equal the Right Min-Heap or exceed it by at most $1$:

8. Complete C++ Implementation for Median Stream

#include <queue>
#include <vector>

class MedianFinder {
private:
    std::priority_queue<int> leftMaxHeap; // Stores smaller half
    std::priority_queue<int, std::vector<int>, std::greater<int>> rightMinHeap; // Stores larger half

public:
    MedianFinder() {}
    
    void addNum(int num) {
        // Step 1: Insert into appropriate heap
        if (leftMaxHeap.empty() || num <= leftMaxHeap.top()) {
            leftMaxHeap.push(num);
        } else {
            rightMinHeap.push(num);
        }
        
        // Step 2: Balance the heaps
        // Left heap can have at most 1 more element than Right heap
        if (leftMaxHeap.size() > rightMinHeap.size() + 1) {
            rightMinHeap.push(leftMaxHeap.top());
            leftMaxHeap.pop();
        } 
        // Right heap cannot exceed Left heap in size
        else if (rightMinHeap.size() > leftMaxHeap.size()) {
            leftMaxHeap.push(rightMinHeap.top());
            rightMinHeap.pop();
        }
    }
    
    double findMedian() {
        if (leftMaxHeap.size() == rightMinHeap.size()) {
            return (double)(leftMaxHeap.top() + rightMinHeap.top()) / 2.0;
        } else {
            // Left heap has one extra element
            return (double)leftMaxHeap.top();
        }
    }
};

9. Summary Comparison of Approaches

10. Summary and Key Takeaways

• **- Exploit Matrix Sorted Properties: Never treat a sorted matrix as raw unsorted data; use row/column sorting invariants to design $\mathcal{O}(n + k \log n)$ heap algorithms.
• **- State Tracking with Pairs: When tracking grid coordinates inside priority queues, wrap values and indices into std::pair structures to let the STL handle primary and secondary sorting automatically.
• **- Dual-Heap Balance Pattern: Finding running medians efficiently relies on splitting datasets into two mirrored halves using a Max-Heap and a Min-Heap.
• **- Heap Size Invariant: Always enforce strict size balancing between heaps (e.g., Left size equals Right size or exceeds by 1) to guarantee $\mathcal{O}(1)$ median retrieval time.
• **- Incremental Optimization Mindset: Always progress from brute-force $\mathcal{O}(n^2 \log n)$ solutions to optimal heap or binary search paradigms during technical interviews.