Heap / Priority Queue
A priority queue is a waiting line where the most important person is always at the front, no matter when they arrived. A heap gives you the min (or max) in O(1) and inserts/removes in O(log n) — perfect for "K largest", "next smallest" and "keep the best so far".
After this topic: You can pick min-heap vs max-heap, keep a heap of size k, and combine two heaps for medians.
Do these first: Trees
Step 1 · Read the lesson
Keep only the best K items; the heap tells you the worst of them in O(log k).
Step 2 · Solve the problems in order
Try each one for about 20 minutes first. Problems with a Run code tab are checked right here. If you are stuck, open Nudge, think again, then Idea. Go skeleton only gives the function shape, and Reference solution is for comparing after you have tried. Tick the box when you could solve it again without help.
- 1.Kth Largest Element in a StreamEasy
You only ever care about the k biggest numbers. Which of them is the "weakest"?
Min-heap of size k. Push each new number; if size exceeds k pop the smallest. The top is the k-th largest.
Target: O(log k) per add
Types such as ListNode, TreeNode and Node are the standard ones for these problems. Define them yourself if you run the code locally.
type KthLargest struct{} func Constructor(k int, nums []int) KthLargest func (k *KthLargest) Add(val int) intTested with go test. Try it yourself first, then compare.
// KthLargest keeps a MIN-heap of the k largest numbers seen. Its root is the smallest of those k, // which is exactly the k-th largest overall. type KthLargest struct { k int h *minInts } type minInts []int func (h minInts) Len() int { return len(h) } func (h minInts) Less(i, j int) bool { return h[i] < h[j] } func (h minInts) Swap(i, j int) { h[i], h[j] = h[j], h[i] } func (h *minInts) Push(x any) { *h = append(*h, x.(int)) } func (h *minInts) Pop() any { old := *h x := old[len(old)-1] *h = old[:len(old)-1] return x } func NewKthLargest(k int, nums []int) *KthLargest { kl := &KthLargest{k: k, h: &minInts{}} for _, n := range nums { kl.Add(n) } return kl } func (kl *KthLargest) Add(val int) int { heap.Push(kl.h, val) if kl.h.Len() > kl.k { heap.Pop(kl.h) // the smallest can no longer be in the top k } return (*kl.h)[0] } - 2.Last Stone WeightEasy
Each round you need the two heaviest stones. Sorting every round is wasteful.
Max-heap: pop two, push back their difference if non-zero, repeat until ≤ 1 stone remains.
Target: O(n log n) time
Types such as ListNode, TreeNode and Node are the standard ones for these problems. Define them yourself if you run the code locally.
func lastStoneWeight(stones []int) intTested with go test. Try it yourself first, then compare.
// LastStoneWeight: each round needs the two heaviest stones, so keep a MAX-heap. // (Go's heap is a min-heap, so store negated weights.) func LastStoneWeight(stones []int) int { h := &minInts{} for _, s := range stones { heap.Push(h, -s) } for h.Len() > 1 { a, b := -heap.Pop(h).(int), -heap.Pop(h).(int) // a >= b if a != b { heap.Push(h, -(a - b)) } } if h.Len() == 0 { return 0 } return -(*h)[0] } - 3.K Closest Points to OriginMedium
You do not need exact distances — comparing squared distances is enough. Keep only the k best.
Max-heap of size k keyed by squared distance; evict the farthest when size > k.
Target: O(n log k) time
Types such as ListNode, TreeNode and Node are the standard ones for these problems. Define them yourself if you run the code locally.
func kClosest(points [][]int, k int) [][]intTested with go test. Try it yourself first, then compare.
// KClosest keeps a MAX-heap of the k closest points so far (keyed by squared distance, no square root needed). // The root is the farthest of the k, the one to evict when a closer point arrives. type pointHeap [][3]int // {squared distance, x, y} func (h pointHeap) Len() int { return len(h) } func (h pointHeap) Less(i, j int) bool { return h[i][0] > h[j][0] } func (h pointHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] } func (h *pointHeap) Push(x any) { *h = append(*h, x.([3]int)) } func (h *pointHeap) Pop() any { old := *h x := old[len(old)-1] *h = old[:len(old)-1] return x } func KClosest(points [][]int, k int) [][]int { h := &pointHeap{} for _, p := range points { heap.Push(h, [3]int{p[0]*p[0] + p[1]*p[1], p[0], p[1]}) if h.Len() > k { heap.Pop(h) } } out := make([][]int, 0, k) for _, e := range *h { out = append(out, []int{e[1], e[2]}) } return out } - 4.Kth Largest Element in an ArrayMedium
Same idea as the stream problem, but all numbers are given at once.
Min-heap of size k (O(n log k)), or quickselect for O(n) average.
Target: O(n log k) time, O(k) space
Types such as ListNode, TreeNode and Node are the standard ones for these problems. Define them yourself if you run the code locally.
func findKthLargest(nums []int, k int) intTested with go test. Try it yourself first, then compare. It is explained step by step on the Heap / Top-K lesson page.
// FindKthLargest: keep a min-heap of the k largest values seen so far. // The root is the smallest of those k, which is exactly the kth largest overall. func FindKthLargest(nums []int, k int) int { h := &IntHeap{} for _, v := range nums { heap.Push(h, v) if h.Len() > k { heap.Pop(h) // evict the smallest: it can no longer be in the top k } } return (*h)[0] } - 5.Task SchedulerMedium
Always do the task with the most remaining runs, but a task must cool down for n intervals after running.
Max-heap of counts plus a queue of (count, time it becomes available again). Each tick run the top task, push it into the cooldown queue, release finished cooldowns. (Math shortcut: (maxCount−1)·(n+1) + number of tasks with maxCount.)
Target: O(T log 26) time
Types such as ListNode, TreeNode and Node are the standard ones for these problems. Define them yourself if you run the code locally.
func leastInterval(tasks []byte, n int) intTested with go test. Try it yourself first, then compare.
// LeastInterval: the most frequent task decides the length. Place it every n+1 slots; the other tasks fill the gaps. // Frames: (maxCount-1) full frames of n+1 slots, plus a last partial row holding every task tied for the max. // If there are so many tasks that no idle time is needed, the answer is simply len(tasks). func LeastInterval(tasks []byte, n int) int { var counts [26]int maxCount := 0 for _, t := range tasks { counts[t-'A']++ maxCount = max(maxCount, counts[t-'A']) } tied := 0 for _, c := range counts { if c == maxCount { tied++ } } return max(len(tasks), (maxCount-1)*(n+1)+tied) } - 6.Design TwitterMedium
getNewsFeed is "merge k sorted lists" in disguise — each followee has a list of tweets by time.
Store tweets per user with a global timestamp and a follow set. For the feed, push each followee's latest tweet into a max-heap by time and pop 10, pushing the previous tweet of that user.
Target: O(k + 10 log k) per feed
Types such as ListNode, TreeNode and Node are the standard ones for these problems. Define them yourself if you run the code locally.
type Twitter struct{} func Constructor() Twitter func (t *Twitter) PostTweet(userId int, tweetId int) func (t *Twitter) GetNewsFeed(userId int) []int func (t *Twitter) Follow(followerId int, followeeId int) func (t *Twitter) Unfollow(followerId int, followeeId int)Tested with go test. Try it yourself first, then compare.
// Twitter keeps each user's tweets (newest last) and follow set. The feed is "merge k sorted lists": // a max-heap on tweet time holds each followed user's newest tweet, and we pop 10, refilling from the same user. type Twitter struct { clock int tweets map[int][]tweet follows map[int]map[int]bool } type tweet struct{ time, id int } func NewTwitter() *Twitter { return &Twitter{tweets: map[int][]tweet{}, follows: map[int]map[int]bool{}} } func (t *Twitter) PostTweet(user, id int) { t.clock++ t.tweets[user] = append(t.tweets[user], tweet{t.clock, id}) } func (t *Twitter) Follow(follower, followee int) { if t.follows[follower] == nil { t.follows[follower] = map[int]bool{} } t.follows[follower][followee] = true } func (t *Twitter) Unfollow(follower, followee int) { delete(t.follows[follower], followee) } type feedItem struct { tw tweet user int index int // position of tw in that user's list, so the previous tweet is index-1 } type feedHeap []feedItem func (h feedHeap) Len() int { return len(h) } func (h feedHeap) Less(i, j int) bool { return h[i].tw.time > h[j].tw.time } func (h feedHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] } func (h *feedHeap) Push(x any) { *h = append(*h, x.(feedItem)) } func (h *feedHeap) Pop() any { old := *h x := old[len(old)-1] *h = old[:len(old)-1] return x } func (t *Twitter) GetNewsFeed(user int) []int { h := &feedHeap{} add := func(u int) { if list := t.tweets[u]; len(list) > 0 { heap.Push(h, feedItem{list[len(list)-1], u, len(list) - 1}) } } add(user) // you see your own tweets too for f := range t.follows[user] { if f != user { add(f) } } feed := []int{} for h.Len() > 0 && len(feed) < 10 { top := heap.Pop(h).(feedItem) feed = append(feed, top.tw.id) if top.index > 0 { heap.Push(h, feedItem{t.tweets[top.user][top.index-1], top.user, top.index - 1}) } } return feed } - 7.Find Median from Data StreamHard
The median sits between the smaller half and the larger half. Which structure gives you the edge of each half cheaply?
Max-heap for the lower half, min-heap for the upper half; keep sizes within 1 of each other. Median is the larger heap's top or the average of both tops.
Target: add O(log n), median O(1)
Types such as ListNode, TreeNode and Node are the standard ones for these problems. Define them yourself if you run the code locally.
type MedianFinder struct{} func Constructor() MedianFinder func (m *MedianFinder) AddNum(num int) func (m *MedianFinder) FindMedian() float64Tested with go test. Try it yourself first, then compare. It is explained step by step on the Heap / Top-K lesson page.
type maxHeap []int func (h maxHeap) Len() int { return len(h) } func (h maxHeap) Less(i, j int) bool { return h[i] > h[j] } func (h maxHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] } func (h *maxHeap) Push(x any) { *h = append(*h, x.(int)) } func (h *maxHeap) Pop() any { old := *h x := old[len(old)-1] *h = old[:len(old)-1] return x } // MedianFinder: a max-heap holds the smaller half, a min-heap the larger half. // Invariant: len(lo) == len(hi) or len(lo) == len(hi)+1, and every lo value <= every hi value. type MedianFinder struct { lo maxHeap // smaller half, root = its largest hi IntHeap // larger half, root = its smallest } func (m *MedianFinder) AddNum(x int) { heap.Push(&m.lo, x) heap.Push(&m.hi, heap.Pop(&m.lo).(int)) // pass the largest of lo across so halves stay ordered if m.hi.Len() > m.lo.Len() { heap.Push(&m.lo, heap.Pop(&m.hi).(int)) // rebalance sizes } } func (m *MedianFinder) FindMedian() float64 { if m.lo.Len() == 0 { return 0 } if m.lo.Len() > m.hi.Len() { return float64(m.lo[0]) } return float64(m.lo[0]+m.hi[0]) / 2 }