题解 | #最小的K个数#

package main

func GetLeastNumbers_Solution(input []int, k int) []int {
	if k == 0 {
		return nil
	}
	// write code here
	heap := newHeap(k)
	for i := 0; i < len(input); i++ {
		heap.Insert(input[i])
	}
	return heap.data[1:]
}

type Heap struct {
	data  []int
	n     int
	count int
}

func newHeap(n int) Heap {
	return Heap{
		data:  make([]int, n+1),
		n:     n,
		count: 0,
	}
}

func (h *Heap) Insert(data int) {
	if h.count >= h.n {
		if h.data[1] > data {
			h.data[1] = data
			h.heapify(1)
		}
		return
	}
	h.count++
	h.data[h.count] = data
	i := h.count
	h.heapify2(i)
}

func (h *Heap) heapify2(i int) {
	for i/2 > 0 && h.data[i] > h.data[i/2] {
		h.data[i], h.data[i/2] = h.data[i/2], h.data[i]
		i /= 2
	}
}

func (h *Heap) heapify(i int) {
	for {
		maxPos := i
		if i*2 <= h.n && h.data[i] < h.data[i*2] {
			maxPos = 2 * i
		}
		if i*2+1 <= h.n && h.data[maxPos] < h.data[2*i+1] {
			maxPos = 2*i + 1
		}
		if maxPos == i {
			break
		}
		h.data[maxPos], h.data[i] = h.data[i], h.data[maxPos]
		i = maxPos
	}
}