题目描述

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

分析

第一种方法是考虑集合中的TreeSet,动态的维护K个最小的数。第二种方法利用快速排序的partition,确定index后，左边的数都比index小，所以当index为k-1时，下标[0,k-1]就是整个数组最小的K个数。

代码一：TreeSet

public ArrayList<Integer> GetLeastNumbers_Solution(int [] input, int k) {
        TreeSet<Integer> ts = new TreeSet<>();
        ArrayList<Integer> list = new ArrayList<>();

        if(k<=0 || k>input.length) return list;
        for(int num:input) {
            if(ts.size()<k) ts.add(num);
            else if(ts.last()>num) {
                ts.pollLast(); //删除最大元素
                ts.add(num);
            }
        } 

        for(Iterator iter = ts.iterator();iter.hasNext();) 
            list.add((Integer)iter.next());
        return list;
    }

方法二：partition

private int partition(int[] arr,int start,int end) {
        if(arr == null || arr.length ==0 || start > end) return -1;
        int pivot = arr[start];
        int i = start;
        int j = end;
        while(i<j) {
            while(i<j && arr[j] > pivot) --j;
            if(i<j) {
                arr[i] = arr[j];
                i++;
            }
            while(i<j && arr[i] < pivot) ++i;
            if(i<j) {
                arr[j] = arr[i];
                --j;
            }
        }
       arr[j] = pivot;
       return j;
    }

    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;
        int index = -1;
        int start = 0;
        int end = input.length-1;
        while(index!=k-1) {
            //System.out.println("start="+start+" end="+end);
            index = partition(input,start,end);
            //System.out.println("index="+index);
            if(index > k-1) end = index-1;
            else start = index+1;
        }
        for(int i=0;i<=index;i++)
            res.add(input[i]);
        return res;
    }