Top K elements

This pattern is useful in scenarios where you need to: Find the K largest or K smallest elements from a collection. Solve problems related to ranking or filtering top results (e.g., top scores, top frequencies, etc.). When sorting, the entire array is unnecessary and inefficient. - Handle a large amount of data or data streams where you cannot store the entire data.

Example Questions:

  • Find the K largest elements in an array.
  • Find the K smallest elements in an array.
  • Find the K most frequent elements in a list.
  • Find the K closest points to the origin.

Common phrases: “Top K”, “K largest”, “K smallest”, “Most frequent”, “Closest K”.

What Are the Variances of It

  1. K Largest Elements: Use a min-heap to track the largest elements while iterating over the array.
  2. K Smallest Elements: Use a max-heap to track the smallest elements.
  3. K Most Frequent Elements: Use a heap with a frequency map to track the most frequent elements.
  4. K Closest Points to Origin: Use a min-heap with distances from the origin.
  5. Data Stream Variants: Maintain a heap of size K while processing a continuous stream of data.

Tricks to Keep in Mind

  • Use Min-Heap for Largest Elements: When finding the K largest elements, keep the K largest elements in a min-heap, and discard the smaller ones.
  • Use Max-Heap for Smallest Elements: For finding the K smallest elements, use a max-heap to track the smallest values.
  • Optimal Time Complexity: This technique typically has time complexity (O(n K)), making it efficient when K is much smaller than n.
  • Sorting isn’t necessary: Sorting the entire array takes (O(n n)), but with a heap, you only care about maintaining the K largest or smallest values.
  • Heap Size: Always maintain the heap size to K during the process, discarding elements as you go.

Common Coding Examples

Question 1: Finding K Largest Elements using Min-Heap

from typing import List
import heapq

class TopKElements:
    def find_k_largest(self, nums: List[int], k: int) -> List[int]:
        """
        Finds the K largest elements in an array using a min-heap.
        """
        # Min-heap of size K
        min_heap = nums[:k]
        heapq.heapify(min_heap)
        
        for num in nums[k:]:
            if num > min_heap[0]:  # If current element is larger than the smallest in the heap
                heapq.heapreplace(min_heap, num)  # Replace the smallest element
        
        return sorted(min_heap, reverse=True)  # Return the K largest in descending order

# Example usage
top_k = TopKElements()
print(top_k.find_k_largest([3, 1, 5, 12, 2, 11], 3))  # Output: [12, 11, 5]

Question 2: Finding K Smallest Elements using Max-Heap

from typing import List
import heapq

class TopKElements:
    def find_k_smallest(self, nums: List[int], k: int) -> List[int]:
        """
        Finds the K smallest elements in an array using a max-heap.
        """
        # Max-heap of size K (by negating values)
        max_heap = [-num for num in nums[:k]]
        heapq.heapify(max_heap)
        
        for num in nums[k:]:
            if -num > max_heap[0]:  # If current element is smaller than the largest in the heap
                heapq.heapreplace(max_heap, -num)  # Replace the largest element
        
        return sorted([-num for num in max_heap])  # Return the K smallest in ascending order

# Example usage
top_k = TopKElements()
print(top_k.find_k_smallest([3, 1, 5, 12, 2, 11], 3))  # Output: [1, 2, 3]

Question 3: Finding K Most Frequent Elements using Min-Heap

from typing import List
import heapq
from collections import Counter

class TopKFrequentElements:
    def find_k_most_frequent(self, nums: List[int], k: int) -> List[int]:
        """
        Finds the K most frequent elements in an array using a min-heap.
        """
        freq_map = Counter(nums)
        
        # Build a min-heap of the K most frequent elements
        min_heap = []
        
        for num, freq in freq_map.items():
            heapq.heappush(min_heap, (freq, num))
            if len(min_heap) > k:
                heapq.heappop(min_heap)
        
        # Extract the numbers from the heap
        return [num for freq, num in min_heap]

# Example usage
top_k_freq = TopKFrequentElements()
print(top_k_freq.find_k_most_frequent([1, 1, 1, 3, 2, 2, 4], 2))  # Output: [1, 2]

Track your progress

Mark this subtopic as completed when you finish reading.