题解 | #最小的K个数#

package main

import "container/heap"

/**
 * 代码中的类名、方法名、参数名已经指定，请勿修改，直接返回方法规定的值即可
 *
 *
 * @param input int整型一维数组
 * @param k int整型
 * @return int整型一维数组
 */
type IntHeap []int

func (h IntHeap) Len() int           { return len(h) }
func (h IntHeap) Less(i, j int) bool { return h[i] > h[j] }
func (h IntHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }

func (h *IntHeap) Push(x interface{}) {
	// Push and Pop use pointer receivers because they modify the slice's length,
	// not just its contents.
	*h = append(*h, x.(int))
}

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

func GetLeastNumbers_Solution(input []int, k int) []int {
	// write code here
	if k == 0 || len(input) == 0 {
		return []int{}
	}

	h := &IntHeap{}
	heap.Init(h)
	for i := 0; i < k; i++ {
		heap.Push(h, input[i])
	}

	for i := k; i < len(input); i++ {
		if input[i] < (*h)[0] {
			heap.Pop(h)
			heap.Push(h, input[i])
		}
	}

	ret := make([]int, k)
	for i := k - 1; i >= 0; i-- {
		ret[i] = heap.Pop(h).(int)
	}

	return ret
}