Heap¶
Table of Contents¶
- 703. Kth Largest Element in a Stream (Easy)
- 1046. Last Stone Weight (Easy)
- 973. K Closest Points to Origin (Medium)
- 215. Kth Largest Element in an Array (Medium)
- 621. Task Scheduler (Medium)
- 355. Design Twitter (Medium)
- 295. Find Median from Data Stream (Hard)
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
1046. Last Stone Weight¶
"""
- Heap
- Time: O(n log n); Space: O(n)
- 0/1 Knapsack
- Time: O(n); Space: O(n)
"""
from heapq import heapify, heappop, heappush
from typing import List
# Heap
def lastStoneWeightHeap(stones: List[int]) -> int:
maxHeap = [-s for s in stones]
heapify(maxHeap)
while len(maxHeap) > 1:
s1 = heappop(maxHeap)
s2 = heappop(maxHeap)
if s1 != s2:
heappush(maxHeap, s1 - s2)
return -maxHeap[0] if maxHeap else 0
# 0/1 Knapsack
def lastStoneWeightKnapsack(stones: List[int]) -> int:
total = sum(stones)
target = total // 2
dp = [0 for _ in range(target + 1)]
for i in stones:
for j in range(target, i - 1, -1):
dp[j] = max(dp[j], dp[j - i] + i)
return total - 2 * dp[target]
if __name__ == "__main__":
stones = [2, 7, 4, 1, 8, 1]
assert lastStoneWeightHeap(stones) == 1
assert lastStoneWeightKnapsack(stones) == 1
#include <iostream>
#include <queue>
#include <vector>
using namespace std;
int lastStoneWeight(vector<int> &stones) {
priority_queue<int> maxHeap(stones.begin(), stones.end());
// Only while there are at least two stones to smash
while (maxHeap.size() > 1) {
int first = maxHeap.top();
maxHeap.pop();
int second = maxHeap.top();
maxHeap.pop();
if (first != second) {
maxHeap.push(first - second);
}
}
return maxHeap.empty() ? 0 : maxHeap.top();
}
int main() {
vector<int> stones = {2, 7, 4, 1, 8, 1};
cout << lastStoneWeight(stones) << endl; // 1
return 0;
}
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]]
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
621. Task Scheduler¶
import heapq
from collections import Counter, deque
from typing import List
def least_interval_heap(tasks: List[str], n: int) -> int:
count = Counter(tasks)
heap = [-c for c in count.values()]
heapq.heapify(heap)
res = 0
q = deque()
while heap or q:
res += 1
if heap:
cnt = 1 + heapq.heappop(heap)
if cnt:
q.append((cnt, res + n))
if q and q[0][1] == res:
heapq.heappush(heap, q.popleft()[0])
return res
def least_interval_math(tasks: List[str], n: int) -> int:
freq = Counter(tasks)
max_exec = max(freq.values())
max_count = sum(1 for v in freq.values() if v == max_exec)
return max((max_exec - 1) * (n + 1) + max_count, len(tasks))
def test_least_interval_heap():
assert least_interval_heap(["A", "A", "A", "B", "B", "B"], 2) == 8
assert least_interval_heap(["A"], 0) == 1
assert least_interval_heap(["A", "A", "A", "B", "B", "B"], 0) == 6
assert least_interval_heap(["A", "A", "A", "A", "B", "B", "B", "C", "C"], 2) == 10
def test_least_interval_math():
assert least_interval_math(["A", "A", "A", "B", "B", "B"], 2) == 8
assert least_interval_math(["A"], 0) == 1
assert least_interval_math(["A", "A", "A", "B", "B", "B"], 0) == 6
assert least_interval_math(["A", "A", "A", "A", "B", "B", "B", "C", "C"], 2) == 10
355. Design Twitter¶
"""
- Similar question: [23. Merge K Sorted Lists](https://leetcode.com/problems/merge-k-sorted-lists/) (Hard)
"""
import heapq
from collections import defaultdict
from typing import List
# Design
class Twitter:
def __init__(self):
self.tweets = defaultdict(list)
self.followees = defaultdict(set)
self.time = 0
def postTweet(self, userId: int, tweetId: int) -> None:
self.tweets[userId].append((self.time, tweetId))
self.time += 1
def getNewsFeed(self, userId: int) -> List[int]:
news_feed = []
news_feed.extend(self.tweets[userId])
for followee in self.followees[userId]:
news_feed.extend(self.tweets[followee])
return [tweet for _, tweet in heapq.nlargest(10, news_feed)]
def follow(self, followerId: int, followeeId: int) -> None:
if followerId != followeeId:
self.followees[followerId].add(followeeId)
def unfollow(self, followerId: int, followeeId: int) -> None:
if followerId != followeeId:
self.followees[followerId].discard(followeeId)
twitter = Twitter()
print(twitter.postTweet(1, 5)) # None
print(twitter.getNewsFeed(1)) # [5]
print(twitter.follow(1, 2)) # None
print(twitter.postTweet(2, 6)) # None
print(twitter.getNewsFeed(1)) # [6, 5]
print(twitter.unfollow(1, 2)) # None
print(twitter.getNewsFeed(1)) # [5]
295. Find Median from Data Stream¶
from heapq import heappop, heappush
class MedianFinder:
def __init__(self):
self.minHeap = []
self.maxHeap = []
self.min_size = 0
self.max_size = 0
def addNum(self, num: int) -> None:
heappush(self.maxHeap, -num)
heappush(self.minHeap, -heappop(self.maxHeap))
self.min_size += 1
if self.min_size > self.max_size:
heappush(self.maxHeap, -heappop(self.minHeap))
self.min_size -= 1
self.max_size += 1
def findMedian(self) -> float:
if self.min_size == self.max_size:
return (-self.maxHeap[0] + self.minHeap[0]) / 2.0
return float(-self.maxHeap[0])
def test_median_finder():
obj = MedianFinder()
obj.addNum(1)
obj.addNum(2)
assert obj.findMedian() == 1.5
obj.addNum(3)
assert obj.findMedian() == 2.0
obj.addNum(4)
assert obj.findMedian() == 2.5
obj.addNum(5)
assert obj.findMedian() == 3.0
#include <iostream>
#include <queue>
#include <vector>
using namespace std;
class MedianFinder {
private:
priority_queue<int> maxHeap;
priority_queue<int, vector<int>, greater<int>> minHeap;
int max_size = 0;
int min_size = 0;
public:
MedianFinder() {}
void addNum(int num) {
if (min_size == max_size) {
minHeap.push(num);
maxHeap.push(minHeap.top());
minHeap.pop();
max_size++;
} else {
maxHeap.push(num);
minHeap.push(maxHeap.top());
maxHeap.pop();
min_size++;
}
}
double findMedian() {
if (min_size == max_size) {
return (maxHeap.top() + minHeap.top()) / 2.0;
} else {
return (double)maxHeap.top();
}
}
};
int main() {
MedianFinder* obj = new MedianFinder();
obj->addNum(1);
obj->addNum(2);
cout << obj->findMedian() << endl; // 1.5
obj->addNum(3);
cout << obj->findMedian() << endl; // 2
return 0;
}