18 · ПРИОРИТЕТ

Куча

container/heap даёт минимум за O(1) и добавление/удаление за O(log n).

2 Go-примитива3 примераInterview 150

СНАЧАЛА ИНСТРУМЕНТЫ

Примитивы Go

Минимальный синтаксис, который понадобится в решении.

Интерфейс

type IntHeap []int
func (h IntHeap) Len() int
func (h IntHeap) Less(i, j int) bool
func (h IntHeap) Swap(i, j int)
func (h *IntHeap) Push(x any)
func (h *IntHeap) Pop() any

Нужно реализовать sort.Interface плюс Push и Pop с pointer receiver.

Операции

heap.Init(&h)
heap.Push(&h, value)
min := heap.Pop(&h).(int)

Для max-heap поменяй Less на h[i] > h[j].

ПРАКТИКА

Примеры решений

Три идиоматичных решения на Go с оценкой времени и памяти.

ПРИМЕР РЕШЕНИЯ

K-й максимальный

Поддерживаем min-heap только из k элементов.

O(n log k) time · O(k) space

solution.go

func findKthLargest(nums []int, k int) int {
    h := &IntHeap{}
    for _, n := range nums {
        heap.Push(h, n)
        if h.Len() > k { heap.Pop(h) }
    }
    return (*h)[0]
}

ПРИМЕР РЕШЕНИЯ

Слияние списков

Куча хранит текущую голову каждого списка.

O(n log k) time · O(k) space

solution.go

func mergeKLists(lists []*ListNode) *ListNode {
    h := &NodeHeap{}
    for _, node := range lists { if node != nil { heap.Push(h, node) } }
    dummy := &ListNode{}; tail := dummy
    for h.Len() > 0 {
        node := heap.Pop(h).(*ListNode)
        tail.Next, tail = node, node
        if node.Next != nil { heap.Push(h, node.Next) }
    }
    return dummy.Next
}

ПРИМЕР РЕШЕНИЯ

K ближайших точек

Куча размера k оставляет только самые близкие точки.

O(n log k) time · O(k) space

solution.go

func kClosest(points [][]int, k int) [][]int {
    h := &PointMaxHeap{}
    for _, point := range points {
        heap.Push(h, point)
        if h.Len() > k { heap.Pop(h) }
    }
    return *h
}

ПЕРЕД КОДОМ

Чек-лист рассуждения

Min-heap или max-heap
Хватит ли кучи размера k
Не забудь type assertion после Pop