跳转至

LC215 - 数组中的第K个最大元素

KthLargestElement.go
package Heap

import (
    "cmp"
    "container/heap"
)

// EasyMinHeap is a generic type.
// The elements should be basic data types.
type EasyMinHeap[T cmp.Ordered] []T

func (h EasyMinHeap[T]) Len() int {
    return len(h)
}

func (h EasyMinHeap[T]) Less(i, j int) bool {
    return h[i] < h[j]
}

func (h EasyMinHeap[T]) Swap(i, j int) {
    h[i], h[j] = h[j], h[i]
}

func (h *EasyMinHeap[T]) Push(x interface{}) {
    *h = append(*h, x.(T))
}

func (h *EasyMinHeap[T]) Pop() interface{} {
    old := *h
    n := len(old)
    x := old[n-1]
    *h = old[0 : n-1]
    return x
}

func (h *EasyMinHeap[T]) Top() interface{} {
    if h.Len() == 0 {
        return nil
    } else {
        return (*h)[0]
    }
}

// 时间复杂度: O(n * log k), 空间复杂度: O(k)
func findKthLargest(nums []int, k int) int {
    minHeap := new(EasyMinHeap[int])
    heap.Init(minHeap)
    for _, num := range nums {
        if minHeap.Len() < k {
            heap.Push(minHeap, num)
        } else if num > minHeap.Top().(int) {
            heap.Push(minHeap, num)
            heap.Pop(minHeap)
        }
    }
    return heap.Pop(minHeap).(int)
}