(11)Go实现的最小堆求前K个最大值

在1,000,000个数字中,选出前100个最大的数字
// 在n个元素中选出前m个元素
// 如果用排序算法,最快时间 NlogN
// 用最小二叉堆形式实现的优先队列,最快时间是 NlogM

算法思路:
(1)最小堆中每次取出来的值都是堆中的最小值,利用这个性质维护一个m个节点最小堆;
(2)遍历一篇数组,每次取出来的值和最小堆中的最小值做对比,如果大于最小堆中的最小值,则和该最小值做替换;
(3)重新获取一次堆中最小值,重复第(2)步,直到遍历数组结束,之后堆中就是前m个最大值;

最小堆的定义和实现
type minHeap struct {
    size int
    nums []int
}

func NewMinHeap() *minHeap {
    return &minHeap{}
}

func Parent(i int) int {
    if i == 0 {
        return 0
    }
    return (i - 1) / 2
}

func LeftChild(i int) int {
    return 2*i + 1
}

func (heap *minHeap) IsEmpty() bool {
    return heap.size == 0
}

func (heap *minHeap) GetSize() int {
    return heap.size
}

func (heap *minHeap) GetMinVal() (int, error) {
    if heap.IsEmpty() {
        return 0, errors.New(
            "failed to geiMinVal,minHeap is empty.")
    }
    return heap.nums[0], nil
}

func siftDown(heap *minHeap, parI int) {
    var minI int
    for {
        leftI := LeftChild(parI)
        switch {
        case leftI+1 > heap.size:
            return
        case leftI+2 > heap.size:
            if heap.nums[parI] > heap.nums[leftI] {
                heap.nums[parI], heap.nums[leftI] = heap.nums[leftI],
                    heap.nums[parI]
            }
            return
        case heap.nums[leftI] <= heap.nums[leftI+1]:
            minI = leftI
        case heap.nums[leftI] > heap.nums[leftI+1]:
            minI = leftI + 1
        }

        if heap.nums[parI] > heap.nums[minI] {
            heap.nums[parI], heap.nums[minI] = heap.nums[minI],
                heap.nums[parI]
        }
        parI = minI
    }
}

func (heap *minHeap) SiftUp(i int) {
    heap.nums = append(heap.nums, i)
    parI := Parent(heap.size)
    childI := heap.size

    for heap.nums[parI] > heap.nums[childI] {
        heap.nums[parI], heap.nums[childI] = heap.nums[childI],
            heap.nums[parI]
        childI = parI
        parI = Parent(parI)
    }
    heap.size++
}

func (heap *minHeap) SiftDown() (int, error) {
    minVal, err := heap.GetMinVal()
    if err != nil {
        return minVal, err
    }

    heap.size--
    heap.nums[0], heap.nums[heap.size] = heap.nums[heap.size], 0

    siftDown(heap, 0)
    return minVal, nil
}

func (heap *minHeap) Replace(i int) (int, error) {
    minVal, err := heap.GetMinVal()
    if err != nil {
        return minVal, err
    }
    heap.nums[0] = i

    siftDown(heap, 0)
    return minVal, nil
}
求前m个最大值的行数定义
func maxNums(nums []int, m int) (resp []int) {
    h := minheap1.NewMinHeap()

    for i := 0; i < m; i++ {
        h.SiftUp(nums[i])
    }

    minVal, _ := h.GetMinVal()
    for _, v := range nums[m:] {
        if v > minVal {
            h.Replace(v)
            minVal, _ = h.GetMinVal()
        }
    }
    for i := 0; i < m; i++ {
        v, _ := h.SiftDown()
        resp = append(resp, v)
    }
    return
}
测试逻辑
func main() {
    buf := []int{}
    for i := 0; i < 1000000; i++ {
        buf = append(buf, 50+rand.Intn(2000000))
    }

    resp := maxNums(buf, 10)
    fmt.Println("最小堆排序出前m个最大值:")
    fmt.Println(resp)

    fmt.Println("----")
    sort.Ints(buf)
    fmt.Println("go自带排序接口排序结果:")
    fmt.Println(buf[len(buf)-10 : len(buf)])
}
测试结果
最小堆排序出前m个最大值:
[2000039 2000039 2000040 2000041 2000041 2000043 2000044 2000048 2000048 2000049]
----
go自带排序接口排序结果:
[2000039 2000039 2000040 2000041 2000041 2000043 2000044 2000048 2000048 2000049]

相关问题:
求n个数字中前m个高频数字:https://www.jianshu.com/p/57918b843a2d

有bug欢迎指出

你可能感兴趣的:((11)Go实现的最小堆求前K个最大值)