最小的K个数

1. 题目描述

输入n个整数，找出其中最小的K个数。例如输入4,5,1,6,2,7,3,8这8个数字，则最小的4个数字是1,2,3,4,。

2. 思路

• 快排
• 大顶堆：求最小的k个数可以使用最大堆，其中堆顶是k个数中的最大数

3. 代码

3.1. 排序

//时间：O(NlogN)
func GetLeastNumbers_Solution(input []int, k int) []int {
	if len(input) == 0 {
		return nil
	}

	if k <= 0 || k > len(input) {
		return nil
	}

	quickSort(input, 0, len(input)-1)
	kNumbers := make([]int, 0)
	for i := 0; i < k; i++ {
		kNumbers = append(kNumbers, input[i])
	}

	return kNumbers
}

func quickSort(data []int, left int, right int) {
	if left >= right {
		return
	}

	p := partition(data, left, right)
	quickSort(data, left, p-1)
	quickSort(data, p+1, right)
}

func partition(data []int, left int, right int) int {
	//生成[l, r]之间的随机索引
	p := left + randInt(0, right-left)
	swap(data, left, p)

	//data[left+1...j] <= v; data[j+1...i] >= v
	i := left + 1
	j := right
	for {
		for i <= j && data[i] < data[left] {
			i++
		}
		for j >= i && data[j] > data[left] {
			j--
		}
		if i >= j {
			break
		}

		swap(data, i, j)
		i++
		j--
	}

	swap(data, left, j)
	return j
}

func swap(data []int, i int, j int) {
	tmp := data[i]
	data[i] = data[j]
	data[j] = tmp
}

func randInt(min, max int) int {
	if min >= max || min == 0 || max == 0 {
		return max
	}
	return rand.Intn(max-min) + min
}

3.2. 堆

• java

public class 最小的K个数
{
    public ArrayList<Integer> GetLeastNumbers_Solution(int[] input, int k)
    {
        ArrayList<Integer> res = new ArrayList<>();
        //检查参数
        if (input == null || input.length == 0 || k <= 0 || k > input.length)
        {
            return res;
        }
        //最小的k个数可以使用最大堆实现，堆顶是最小的k个数中的最大数
        PriorityQueue<Integer> maxHeap = new PriorityQueue<>(k, new Comparator<Integer>()
        {
            @Override
            public int compare(Integer o1, Integer o2)
            {
                return o2 - o1;
            }
        });
        //遍历最大堆插入到list中
        for (int val : input)
        {
            //如果数量还没达到k，那么插入到最大堆中
            if (maxHeap.size() < k)
            {
                maxHeap.offer(val);
            }
            else
            {
                //如果val比堆顶小，那么删除堆顶，再插入这个数
                if (val < maxHeap.peek())
                {
                    maxHeap.poll();
                    maxHeap.offer(val);
                }
            }
        }

        for (Integer val : maxHeap)
        {
            res.add(val);
        }

        return res;
    }
}

• go

import (
	"container/heap"
	"sort"
)

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) Pop() interface{} {
	old := *h
	n := len(old)
	x := old[n-1]
	*h = old[:n-1]
	return x
}

func (h *IntHeap) Push(v interface{}) {
	*h = append(*h, v.(int))
}

func GetLeastNumbers_Solution2(input []int, k int) []int {
	// write code here
	if k <= 0 || k > len(input) {
		return []int{}
	}
	//     sort.Ints(input)
	//     return input[:k]
	h := IntHeap{}
	heap.Init(&h)
	for _, v := range input {
		if h.Len() < k {
			heap.Push(&h, v)
		} else if h.Len() == k {
			temp := heap.Pop(&h).(int)
			if temp > v {
				heap.Push(&h, v)
			} else {
				heap.Push(&h, temp)
			}
		}
	}
	sort.Ints(h)
	return h
}

4. 参考

• 最小的K个数_牛客题霸_牛客网
• 面试题 17.14. 最小K个数 - 力扣（LeetCode）