Heap Top K¶
Table of Contents¶
- 215. Kth Largest Element in an Array (Medium)
- 973. K Closest Points to Origin (Medium)
- 347. Top K Frequent Elements (Medium)
- 692. Top K Frequent Words (Medium)
- 264. Ugly Number II (Medium)
- 451. Sort Characters By Frequency (Medium)
- 703. Kth Largest Element in a Stream (Easy)
- 767. Reorganize String (Medium)
- 786. K-th Smallest Prime Fraction (Medium)
215. Kth Largest Element in an Array¶
import heapq
import random
from typing import List
class Solution:
@staticmethod
def _partition(nums: List[int], left: int, right: int) -> tuple[int, int]:
"""Split a range into values smaller than, equal to, and greater than pivot."""
pivot = nums[(left + right) // 2]
smaller = current = left
greater = right
while current <= greater:
if nums[current] < pivot:
nums[smaller], nums[current] = nums[current], nums[smaller]
smaller += 1
current += 1
elif nums[current] > pivot:
nums[current], nums[greater] = nums[greater], nums[current]
greater -= 1
else:
current += 1
return smaller, greater
def sort(self, nums: List[int], k: int) -> int:
"""Built-in sort: O(n log n) time, O(n) space in Python."""
nums.sort(reverse=True)
return nums[k - 1]
def findKthLargest(self, nums: List[int], k: int) -> int:
"""LeetCode entry point using randomized in-place quickselect."""
return self.quick_select_in_place(nums, k)
def min_heap(self, nums: List[int], k: int) -> int:
"""Min-heap: O(n log k) time, O(k) space."""
min_heap = []
for num in nums:
heapq.heappush(min_heap, num)
if len(min_heap) > k:
heapq.heappop(min_heap)
return min_heap[0]
def quick_sort(self, nums: List[int], k: int) -> int:
"""In-place quicksort: O(n log n) average time, O(log n) average space."""
ranges = [(0, len(nums) - 1)]
while ranges:
left, right = ranges.pop()
if left >= right:
continue
equal_start, equal_end = self._partition(nums, left, right)
left_range = (left, equal_start - 1)
right_range = (equal_end + 1, right)
# Process the smaller range first to keep the stack small.
if left_range[1] - left_range[0] > right_range[1] - right_range[0]:
left_range, right_range = right_range, left_range
ranges.append(right_range)
ranges.append(left_range)
return nums[-k]
def quick_select_recursive(self, nums: List[int], k: int) -> int:
"""
Find the kth largest using recursive list partitions.
k largest
3, 2, 1, 5, 6, 4
0 1 2 3 4 5
Time: O(n) average, O(n^2) worst case.
Space: O(n) average for partitions and recursion.
"""
pivot = random.choice(nums)
left = [number for number in nums if number > pivot]
mid = [number for number in nums if number == pivot]
right = [number for number in nums if number < pivot]
if len(left) >= k:
return self.quick_select_recursive(left, k)
if len(left) + len(mid) >= k:
return pivot
return self.quick_select_recursive(right, k - len(left) - len(mid))
def quick_select_in_place(self, nums: List[int], k: int) -> int:
"""
Find the kth largest using randomized in-place quickselect.
The kth largest is the (n - k)th smallest using a zero-based index.
Time: O(n) average, O(n^2) worst case.
Space: O(1), because partitioning is iterative and in place.
"""
target = len(nums) - k
def partition(left: int, right: int) -> int:
pivot = random.randint(left, right)
nums[pivot], nums[right] = nums[right], nums[pivot]
pivot = nums[right]
store = left
for index in range(left, right):
if nums[index] < pivot:
nums[index], nums[store] = nums[store], nums[index]
store += 1
nums[store], nums[right] = nums[right], nums[store]
return store
left, right = 0, len(nums) - 1
while left <= right:
pivot_index = partition(left, right)
if pivot_index == target:
return nums[pivot_index]
if pivot_index < target:
left = pivot_index + 1
else:
right = pivot_index - 1
raise RuntimeError("quickselect failed to find the target")
if __name__ == "__main__":
solution = Solution()
methods = (
solution.findKthLargest,
solution.sort,
solution.min_heap,
solution.quick_sort,
solution.quick_select_recursive,
solution.quick_select_in_place,
)
cases = (
([3, 2, 1, 5, 6, 4], 2, 5),
([3, 2, 3, 1, 2, 4, 5, 5, 6], 4, 4),
([-1, -1], 2, -1),
([1], 1, 1),
)
for method in methods:
for nums, k, expected in cases:
assert method(nums.copy(), k) == expected
973. K Closest Points to Origin¶
-
Tags: Array, Math, Divide And Conquer, Geometry, Sorting, Heap Priority Queue, Quickselect
import heapq
from typing import List
def kClosest(points: List[List[int]], k: int) -> List[List[int]]:
heap = []
for x, y in points:
dist = -(x**2 + y**2) # max heap
if len(heap) < k:
heapq.heappush(heap, (dist, x, y))
else:
heapq.heappushpop(heap, (dist, x, y)) # push and pop the smallest
return [[x, y] for (_, x, y) in heap]
# Time complexity: O(n * log(k))
# - O(log(k)) for heapify
# - O(n) for iterating through the input list
# Space complexity: O(k)
points = [[1, 3], [-2, 2]]
k = 1
print(kClosest(points, k)) # [[-2, 2]]
347. Top K Frequent Elements¶
-
Tags: Array, Hash Table, Divide And Conquer, Sorting, Heap Priority Queue, Bucket Sort, Counting, Quickselect
import heapq
from collections import Counter
from typing import List
# Heap + Counter
def topKFrequent(nums: List[int], k: int) -> List[int]:
minHeap = []
for val, freq in Counter(nums).items():
if len(minHeap) < k:
heapq.heappush(minHeap, (freq, val))
else:
heapq.heappushpop(minHeap, (freq, val))
return [i for _, i in minHeap]
# Counter (Most Common)
def topKFrequentCounter(nums: List[int], k: int) -> List[int]:
commons = Counter(nums).most_common(k)
return [i for i, _ in commons]
nums = [1, 1, 1, 2, 2, 3]
k = 2
print(topKFrequent(nums, k)) # [1, 2]
print(topKFrequentCounter(nums, k)) # [1, 2]
692. Top K Frequent Words¶
-
Tags: Array, Hash Table, String, Trie, Sorting, Heap Priority Queue, Bucket Sort, Counting
import heapq
from collections import Counter
from typing import List
class WordFrequency:
def __init__(self, word, freq):
self.word = word
self.freq = freq
def __lt__(self, other):
# If the frequency is different
if self.freq != other.freq:
# The word with the lower frequency comes first
return self.freq < other.freq
else:
# The word with the lower alphabetical order comes first
return self.word > other.word
def topKFrequent(words: List[str], k: int) -> List[str]:
heap = []
for word, freq in Counter(words).items():
heapq.heappush(heap, WordFrequency(word, freq))
if len(heap) > k:
heapq.heappop(heap)
heap.sort(reverse=True)
return [x.word for x in heap]
words = ["i", "love", "leetcode", "i", "love", "coding"]
k = 2
print(topKFrequent(words, k)) # ["i", "love"]
264. Ugly Number II¶
import heapq
def nthUglyNumber(n: int) -> int:
heap = [1]
seen = set(heap)
factors = [2, 3, 5]
current = 1
# Pop the smallest ugly number n times
for _ in range(n):
current = heapq.heappop(heap) # Pop the smallest ugly number
for factor in factors:
new = current * factor
if new not in seen:
seen.add(new)
heapq.heappush(heap, new)
return current
print(nthUglyNumber(10)) # 12
451. Sort Characters By Frequency¶
703. Kth Largest Element in a Stream¶
import heapq
from typing import List
class KthLargest:
def __init__(self, k: int, nums: List[int]):
"""Min Heap Setup: O(n log k) time, O(k) space.
Keeping only k values means the smallest heap item is the kth largest.
"""
self.k = k
self.heap = []
for num in nums:
self.add(num)
def add(self, val: int) -> int:
"""Min Heap Stream Update: O(log k) time, O(k) space.
Push each value and remove the smallest once more than k candidates exist.
"""
heapq.heappush(self.heap, val)
if len(self.heap) > self.k:
heapq.heappop(self.heap)
return self.heap[0]
def test_kth_largest():
obj = KthLargest(3, [4, 5, 8, 2])
assert obj.add(3) == 4
assert obj.add(5) == 5
assert obj.add(10) == 5
assert obj.add(9) == 8
assert obj.add(4) == 8
obj = KthLargest(1, [])
assert obj.add(-3) == -3
assert obj.add(-2) == -2
assert obj.add(-4) == -2
obj = KthLargest(2, [0])
assert obj.add(-1) == -1
assert obj.add(1) == 0
assert obj.add(1) == 1
767. Reorganize String¶
import heapq
from collections import Counter
def reorganizeString(s: str) -> str:
if not s:
return ""
heap = [(-freq, char) for char, freq in Counter(s).items()] # max heap
heapq.heapify(heap)
result = []
prev_count, prev_char = 0, ""
while heap:
count, char = heapq.heappop(heap) # pop the most frequent character
result.append(char) # append the character to the result
if prev_count < 0: # if the previous character still has remaining count
heapq.heappush(heap, (prev_count, prev_char))
prev_count = count + 1 # update the current character's remaining count
prev_char = char # update the current character
# check if there is any invalid result
if len(result) != len(s):
return ""
return "".join(result)
print(reorganizeString("aab"))
786. K-th Smallest Prime Fraction¶
import heapq
from typing import List
def kthSmallestPrimeFraction(arr: List[int], k: int) -> List[int]:
max_heap = []
for j in range(1, len(arr)):
for i in range(j):
heapq.heappush(max_heap, (-arr[i] / arr[j], arr[i], arr[j]))
if len(max_heap) > k:
heapq.heappop(max_heap)
return [max_heap[0][1], max_heap[0][2]]
arr = [1, 2, 3, 5]
k = 3
print(kthSmallestPrimeFraction(arr, k)) # [2, 5]