剑指offer（31）：最小的k个数

题目描述

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

分析

最简单直接的思路：调用java类库函数进行排序（时间复杂度为$O(nlogn)$，取前k个数即可。

牛客AC代码：

import java.util.*;

public class Solution {
    public ArrayList<Integer> GetLeastNumbers_Solution(int [] input, int k) {
        ArrayList<Integer> list = new ArrayList<Integer>();
        if(input == null ||
            input.length <=0 ||
            k > input.length ||
            k <= 0)
            return list;

        int length = input.length;
        for(int i = 0; i < length; i++)
            list.add(input[i]);

        Collections.sort(list);

        while(list.size() > k) {
            list.remove(k);
        }

        return list;
    }
}

推荐思路

当内存有限制而数组非常大是，是无法实现直接排序的。

创建一个用于保存最小k个数的容器（TreeSet，默认从小到大排序），第一步是放入k个元素，第二步是比较容器中最大的数与当前数组遍历的那个数的大小，如果容器中最大的数比当前遍历的数大，就移除该数，并放入当前遍历的那个数。这样整个循环结束的时候，容器中就是最小的k个数了。

$O(nlogn)$的算法，适合处理海量数据（BAT）。

牛客AC代码：

import java.util.*;

public class Solution {
    public ArrayList<Integer> GetLeastNumbers_Solution(int [] input, int k) {
        ArrayList<Integer> list = new ArrayList<Integer>();
        if(input == null || k <= 0 || input.length == 0 || k > input.length) 
            return list;

        //创建一TreeSet集合，因为TreeSet默认就是从小到大排序的
        TreeSet<Integer> set = new TreeSet<Integer>();
        for (int i = 0; i < input.length; i++) {
            if(set.size() < k){
                set.add(input[i]);
            }else{
                int maxVal = set.last();
                if(input[i] < maxVal) {
                    set.remove(maxVal);
                    set.add(input[i]);
                }
            }
        }

        //循环结束，s中已经保存了前k个最小的数，只需要迭代放入ArrayList中即可
        Iterator<Integer> it = set.iterator();
        while(it.hasNext()){
            list.add(it.next());
        }
        return list;
    }
}

参考
1. 何海涛，剑指offer名企面试官精讲典型编程题（纪念版），电子工业出版社