Golang算法模板-优先队列

215.数组中的第K个最大元素

https://leetcode.cn/problems/kth-largest-element-in-an-array/description/?envType=study-plan-v2&envId=leetcode-75

题目描述

给定整数数组 nums 和整数 k，请返回数组中第 **k** 个最大的元素。

请注意，你需要找的是数组排序后的第 k 个最大的元素，而不是第 k 个不同的元素。

你必须设计并实现时间复杂度为 O(n) 的算法解决此问题。

解题思路：使用小根堆，容量只有k个，每次取出最小的元素，和数组中下一个元素比较，如果下一个元素更大，则下一个元素入队，否则再让刚才取出的元素入队，这样遍历完数组之后，队首就是小根堆的根，及k个最大的值中最小的元素了

优先队列底层使用堆来完成，小根堆或大根堆

import (
    "container/heap"    //内置包
    "fmt"
)
​
type Stru struct {
    num int
}
type StruArr []Stru
​
func (struArr *StruArr) Len() int {
    return len(*struArr)
}
func (struArr *StruArr) Less(i, j int) bool {
    return (*struArr)[i].num < (*struArr)[j].num //小根堆
}
func (struArr *StruArr) Swap(i, j int) {
    (*struArr)[i], (*struArr)[j] = (*struArr)[j], (*struArr)[i]
}
func (struArr *StruArr) Push(x interface{}) { //针对切片压入最后一个元素
    *struArr = append(*struArr, x.(Stru))
}
func (struArr *StruArr) Pop() interface{} { //针对切片弹出最后一个元素
    l := len(*struArr)
    x := (*struArr)[l-1]
    *struArr = (*struArr)[:l-1]
    return x
}
​
func findKthLargest(nums []int, k int) int {
    struArr := &StruArr{}
    heap.Init(struArr)
    for i := 0; i < k; i++ {
        stru := Stru{num: nums[i]}
        heap.Push(struArr, stru)
    }
    for i := k; i < len(nums); i++ {
        t := heap.Pop(struArr).(Stru)
        if nums[i] > t.num {
            stru := Stru{num: nums[i]}
            heap.Push(struArr, stru)
        } else {
            heap.Push(struArr, t)
        }
    }
​
    return heap.Pop(struArr).(Stru).num
}