Java 解决剑指offer 最小的k个数

题目描述

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

题目解析

这道题最简单的思路是把数组排序，排序的时间复杂度为O(NlogN)。但是我们有更好得解决方案。

1.时间复杂度为O(n)，但需要改变数组。

我们将思路放到快速排序上，在快速排序的Partition中，我们每次选择一个元素，将数组中小于这个元素的值放在这个元素的左边，不小于这个元素的值放在这个元素的右边，并返回这个元素所在的位置。也就是说，我们如果调用Partition方法的返回值是k-1，则数组中0到k-1位置的k个元素即为我们要找的元素。如果不是，则根据返回值的位置继续进行Partition操作。
代码如下：

import java.util.*;
public class Solution {
    public ArrayList<Integer> GetLeastNumbers_Solution(int [] input, int k) {
        ArrayList<Integer> list = new ArrayList();
        if(input==null||input.length<k||k<=0){
            return list;
        }
        int lo = 0;
        int hi = input.length-1;
        int index = partition(input,lo,hi);
        while(index+1!=k){
            if(index<k){
                lo = index+1;
                index = partition(input,lo,hi);
            }else{
                hi = index-1;
                index = partition(input,lo,hi);
            }
        }
        for(int i=0;i<k;i++){
            list.add(input[i]);
        }
        return list;
    }
    private int partition(int[] input,int lo,int hi){
        if(lo==hi){
            return lo;
        }
        int temp = input[lo];
        int i=lo;
        int j=hi;
        while(i<j){
            while(input[j]>=temp&&j>i){
                j--;
            }
            input[i] = input[j];
            while(input[i]<temp&&i<j){
                i++;
            }
            if(i<j){
                input[j] = input[i];
            }
        }
        input[j] = temp;
        return j;
    }
}

上述方法的明显不足就是我们需要变动输入的数组。

2.时间复杂度为O(Nlogk)的方法，特别适合海量数据。
另一个思路就是我们创建一个容量为k的容器，用来保存最小的k个元素。再遍历数组时，当前元素如果小于容器中的最大值，则将其替换为当前数字，直到遍历结束剩下的元素就是最小的k个元素。容器的选择我们的第一反应应该是大根堆，大根堆使我们每次可以在O(1)的时间复杂度内找到容器中的最大元素，在O(logk)的时间复杂度内完成元素的删除和插入。
下面是基于大根堆的代码：

import java.util.*;
public class Solution {
    private int[] container;
    public ArrayList<Integer> GetLeastNumbers_Solution(int [] input, int k) {
        ArrayList<Integer> list = new ArrayList();
        if(input==null||input.length<k||k<=0){
            return list;
        }
        container = new int[k+1];
        for(int i=0;i<k;i++){
            container[i+1] = input[i];
            swim(i+1);
        }
        for(int i=k;i<input.length;i++){
            if(input[i]<container[1]){
                container[1] = input[i];
                sink(1);
            }
        }
        for(int i=1;i<k+1;i++){
            list.add(container[i]);
        }
        return list;
    }
    private void sink(int k){
        while((k<<1)<container.length){
            int j = (k<<1);
            if(j+1<container.length&&container[j]<container[j+1]){
                j++;
            }
            if(container[k]>=container[j]){
                return;
            }
            swap(k,j);
            k = j;
        }
    }
    private void swim(int k){
        while(k>1&&container[k>>1]<container[k]){
            swap(k>>1,k);
            k = (k>>1);
        }
    }
    private void swap(int i,int j){
        int t = container[i];
        container[i] = container[j];
        container[j] = t;
    }
}

除了大根堆之外，查找、删除、插入复杂度为O(logk)的容器还有红黑树，当然，红黑树的实现过于复杂，如果我们自己实现一颗红黑树，无论是在面试、考试还是比赛中，都是不现实的。但是Java中为我们提供了一个基于红黑树的集合TreeSet，如果允许的话我们可以直接使用。
下面是基于TreeSet的代码：

import java.util.*;
public class Solution {
    public ArrayList<Integer> GetLeastNumbers_Solution(int [] input, int k) {
        ArrayList<Integer> list = new ArrayList();
        TreeSet<Integer> set = new TreeSet();
        if(input==null||input.length<k||k<=0){
            return list;
        }
        for(int i=0;i<k;i++){
            set.add(input[i]);
        }
        for(int i=k;i<input.length;i++){
            if(input[i]<set.last()){
                set.pollLast();
                set.add(input[i]);
            }
        }
        for(int i:set){
            list.add(i);
        }
        return list;
    }
}