Coding Patterns: Top K Numbers
In Coding Patterns series, we will try to recognize common patterns underlying behind each algorithm question, using real examples from Leetcode.
Previous posts were about Sliding Window, Two Pointers, Fast & Slow Pointers, Merge Intervals, Cyclic Sort, In-place Reversal of a Linked List, Breadth First Search (BFS), Depth First Search (DFS), Two Heaps, Subsets and Modified Binary Search patterns and today, we will introduce Top K Numbers pattern which is very useful to solve the problems that asks us to find the top / smallest / frequent K elements among a given set.
We are going to use Heap data structure to keep track of K elements.
Problem: K-th Largest Element
LeetCode 215 - Kth Largest Element in an Array [medium]
Find the kth largest element in an unsorted array. Note that it is the kth largest element in the sorted order, not the kth distinct element.
Example 1:
Input: [3, 2, 1, 5, 6, 4] and k = 2
Output: 5
Example 2:
Input: [3, 2, 3, 1, 2, 4, 5, 5, 6] and k = 4
Output: 4
Note:
You may assume k is always valid, 1 ≤ k ≤ array’s length.
Top K Numbers Solution
The best data structure to keep track of top K elements is Heap.
If we iterate through the array one element at a time and keep kth largest element in a heap such that each time we find a larger number than the smallest number in the heap, we do two things:
- Take out the smallest number from the heap
- Insert the larger number into the heap
This will ensure that we always have top k largest numbers in the heap. We will use a min-heap for this;
from heapq import *
class Solution:
def findKthLargest(self, nums: List[int], k: int) -> int:
min_heap = []
for i in range(k):
heappush(min_heap, nums[i])
for i in range(k, len(nums)):
if nums[i] > min_heap[0]:
heappop(min_heap)
heappush(min_heap, nums[i])
return min_heap[0]
Time Complexity: O(N log K).
Space Complexity: O(K)
How to identify?
If the problem asking us to find the top / smallest / frequent K elements among a given set, we need to think about Top K Numbers pattern.
While solving the problems, we are going to use Heap data structure to keep track of K elements.