Common Heap Patterns & Problem Types

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44

import heapq
from typing import List
 
def find_kth_largest(nums: List[int], k: int) -> int:
    """
    Find the k-th largest element using a min-heap of size k.
    
    Time Complexity: O(n log k)
    Space Complexity: O(k)
    
    The heap maintains the k largest elements seen so far.
    Its root (minimum) is the k-th largest overall.
    """
    # Python's heapq is a min-heap by default - perfect for this problem!
    min_heap = []
    
    for num in nums:
        if len(min_heap) < k:
            # Phase 1: Fill the heap with first k elements
            heapq.heappush(min_heap, num)
        elif num > min_heap[0]:
            # Phase 2: Replace minimum if current is larger
            # heapreplace is more efficient than pop + push
            heapq.heapreplace(min_heap, num)
        # If num <= min_heap[0], skip it - can't be in top k
    
    # The root of min-heap is the k-th largest
    return min_heap[0]
 
 
# Example usage
nums = [3, 2, 1, 5, 6, 4]
k = 2
result = find_kth_largest(nums, k)
print(f"The {k}th largest element is: {result}")  # Output: 5
 
# Trace through the algorithm:
# Process 3: heap = [3]
# Process 2: heap = [2, 3]  (k=2 elements now)
# Process 1: 1 < 2, skip
# Process 5: 5 > 2, replace: heap = [3, 5]
# Process 6: 6 > 3, replace: heap = [5, 6]
# Process 4: 4 < 5, skip
# Final heap minimum = 5 = 2nd largest ✓

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41

import heapq
from typing import List
 
def find_kth_smallest(nums: List[int], k: int) -> int:
    """
    Find the k-th smallest element using a max-heap of size k.
    
    Python's heapq is a min-heap, so we negate values to simulate max-heap.
    
    Time Complexity: O(n log k)
    Space Complexity: O(k)
    """
    max_heap = []  # Will store negated values
    
    for num in nums:
        if len(max_heap) < k:
            # Phase 1: Fill the heap with first k elements (negated)
            heapq.heappush(max_heap, -num)
        elif num < -max_heap[0]:  # num < current maximum
            # Phase 2: Replace maximum if current is smaller
            heapq.heapreplace(max_heap, -num)
        # If num >= -max_heap[0], skip it - can't be in bottom k
    
    # The root of max-heap (negated) is the k-th smallest
    return -max_heap[0]
 
 
# Example usage
nums = [3, 2, 1, 5, 6, 4]
k = 2
result = find_kth_smallest(nums, k)
print(f"The {k}th smallest element is: {result}")  # Output: 2
 
# Trace:
# Process 3: heap = [-3] → represents max-heap [3]
# Process 2: heap = [-3, -2] → max-heap [3, 2]
# Process 1: 1 < 3, replace: heap = [-2, -1] → max-heap [2, 1]
# Process 5: 5 > 2, skip
# Process 6: 6 > 2, skip
# Process 4: 4 > 2, skip
# Final: -heap[0] = 2 = 2nd smallest ✓

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42

import heapq
from typing import List
 
def find_kth_largest_optimized(nums: List[int], k: int) -> int:
    """
    Optimized version that uses the smaller heap.
    
    If k > n/2, it's more efficient to find (n-k+1)th smallest
    using a max-heap of size (n-k+1).
    
    This ensures heap size is at most ceil(n/2), 
    reducing worst-case operations.
    """
    n = len(nums)
    
    # Use the more efficient formulation
    if k > n // 2:
        # Find (n-k+1)th smallest instead
        k_small = n - k + 1
        max_heap = []
        for num in nums:
            if len(max_heap) < k_small:
                heapq.heappush(max_heap, -num)
            elif num < -max_heap[0]:
                heapq.heapreplace(max_heap, -num)
        return -max_heap[0]
    else:
        # Standard min-heap approach for k-th largest
        min_heap = []
        for num in nums:
            if len(min_heap) < k:
                heapq.heappush(min_heap, num)
            elif num > min_heap[0]:
                heapq.heapreplace(min_heap, num)
        return min_heap[0]
 
 
# Example: Finding 999,999th largest in 1,000,000 elements
# Without optimization: heap of 999,999 elements
# With optimization: heap of 2 elements (finding 2nd smallest)
nums = list(range(1_000_000))
print(find_kth_largest_optimized(nums, 999_999))  # Much faster!

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65

import heapq
from typing import Optional
 
class KthLargestStream:
    """
    Maintains the k-th largest element in a stream.
    
    Supports:
    - add(val): O(log k) - Add a new element
    - get_kth(): O(1) - Get current k-th largest
    
    Use case: Real-time leaderboard, percentile tracking, 
              top-k monitoring
    """
    
    def __init__(self, k: int, initial_nums: list[int] = None):
        """
        Initialize with k and optionally some initial numbers.
        """
        self.k = k
        self.min_heap = []
        
        if initial_nums:
            for num in initial_nums:
                self.add(num)
    
    def add(self, val: int) -> Optional[int]:
        """
        Add a new value to the stream.
        Returns the current k-th largest (or None if < k elements).
        """
        if len(self.min_heap) < self.k:
            heapq.heappush(self.min_heap, val)
        elif val > self.min_heap[0]:
            heapq.heapreplace(self.min_heap, val)
        
        return self.get_kth()
    
    def get_kth(self) -> Optional[int]:
        """
        Get the current k-th largest.
        Returns None if fewer than k elements have been seen.
        """
        if len(self.min_heap) < self.k:
            return None
        return self.min_heap[0]
 
 
# Example: Real-time stock price monitoring
# Track the 3rd highest price seen throughout the day
tracker = KthLargestStream(k=3)
 
prices = [105.2, 102.8, 108.5, 101.3, 110.0, 107.2, 112.5]
for price in prices:
    kth = tracker.add(price)
    print(f"New price: {price:>6.1f} | 3rd highest so far: {kth}")
 
# Output:
# New price:  105.2 | 3rd highest so far: None
# New price:  102.8 | 3rd highest so far: None
# New price:  108.5 | 3rd highest so far: 102.8
# New price:  101.3 | 3rd highest so far: 102.8
# New price:  110.0 | 3rd highest so far: 105.2
# New price:  107.2 | 3rd highest so far: 107.2
# New price:  112.5 | 3rd highest so far: 108.5

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55

import random
from typing import List
 
def quickselect(nums: List[int], k: int) -> int:
    """
    Find k-th largest using QuickSelect algorithm.
    
    Time: O(n) average, O(n²) worst case
    Space: O(1) extra (modifies input)
    
    Note: Modifies the input array!
    """
    # Convert to 0-indexed position from end
    target_idx = len(nums) - k
    
    def partition(left: int, right: int, pivot_idx: int) -> int:
        """Partition around pivot, return final pivot position."""
        pivot = nums[pivot_idx]
        # Move pivot to end
        nums[pivot_idx], nums[right] = nums[right], nums[pivot_idx]
        
        store_idx = left
        for i in range(left, right):
            if nums[i] < pivot:
                nums[store_idx], nums[i] = nums[i], nums[store_idx]
                store_idx += 1
        
        # Move pivot to final position
        nums[right], nums[store_idx] = nums[store_idx], nums[right]
        return store_idx
    
    def select(left: int, right: int) -> int:
        """Recursively select until target found."""
        if left == right:
            return nums[left]
        
        # Random pivot for O(n) expected time
        pivot_idx = random.randint(left, right)
        pivot_idx = partition(left, right, pivot_idx)
        
        if target_idx == pivot_idx:
            return nums[pivot_idx]
        elif target_idx < pivot_idx:
            return select(left, pivot_idx - 1)
        else:
            return select(pivot_idx + 1, right)
    
    return select(0, len(nums) - 1)
 
 
# Example
nums = [3, 2, 1, 5, 6, 4]
k = 2
result = quickselect(nums.copy(), k)  # Copy to preserve original
print(f"The {k}th largest: {result}")  # 5

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59

import heapq
from collections import Counter
from typing import List
 
def top_k_frequent(nums: List[int], k: int) -> List[int]:
    """
    LeetCode 347: Top K Frequent Elements
    
    Return the k most frequent elements.
    Uses min-heap of size k, keyed by frequency.
    
    Time: O(n log k)
    Space: O(n) for counter + O(k) for heap
    """
    # Count frequencies
    freq = Counter(nums)
    
    # Min-heap storing (frequency, element)
    # We want k highest frequencies, so use min-heap of size k
    min_heap = []
    
    for num, count in freq.items():
        if len(min_heap) < k:
            heapq.heappush(min_heap, (count, num))
        elif count > min_heap[0][0]:
            heapq.heapreplace(min_heap, (count, num))
    
    # Extract elements (not frequencies)
    return [num for count, num in min_heap]
 
 
def k_closest_points(points: List[List[int]], k: int) -> List[List[int]]:
    """
    LeetCode 973: K Closest Points to Origin
    
    Return k points closest to origin (0, 0).
    Uses max-heap of size k, keyed by distance.
    
    Time: O(n log k)
    Space: O(k)
    """
    # Max-heap storing (-distance_squared, point)
    # We want k smallest distances, so use max-heap of size k
    max_heap = []
    
    for x, y in points:
        dist_sq = x*x + y*y
        
        if len(max_heap) < k:
            heapq.heappush(max_heap, (-dist_sq, [x, y]))
        elif dist_sq < -max_heap[0][0]:
            heapq.heapreplace(max_heap, (-dist_sq, [x, y]))
    
    return [point for dist, point in max_heap]
 
 
# Examples
print(top_k_frequent([1,1,1,2,2,3], 2))  # [1, 2]
print(k_closest_points([[1,3],[-2,2]], 1))  # [[-2, 2]]