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;
    }
}