【剑指Offer-时间效率】面试题40：最小的k个数

题目描述

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

思路1

可以先将这n个整数升序排序，然后输出前k个数字即可。下面的代码使用快速排序：

class Solution {
public:
    
    void swap(vector<int>& v, int i, int j){
        int t = v[i];
        v[i] = v[j];
        v[j] = t;
    }
    
    int partition(vector<int>& v, int left, int right){
        int i = left;
        int j = right;
        int baseline = v[left];
        while(i<j){
            while(v[j]>=baseline && i<j)    //先找合适的j
                j--;
            while(v[i]<=baseline && i<j)    //再找合适的i
                i++;
            swap(v, i, j);
        }
        swap(v, left, i);
        return i;
    }
    
    void quicksort(vector<int>& v, int left, int right){
        if(left<right){
            int idx = partition(v, left, right);
            quicksort(v, left, idx-1);
            quicksort(v, idx+1, right);
        }
    }
    
    vector<int> GetLeastNumbers_Solution(vector<int> input, int k) {
        vector<int> ans;
        if(input.empty() || k<=0 || k>input.size())
            return ans;
        
        quicksort(input, 0, input.size()-1);
        for(int i=0; i<k&&i<input.size(); i++)
            ans.push_back(input[i]);
        return ans;
    }
};

由于使用快速排序，所以时间复杂度为O(nlogn).

思路2

第二种方法类似于面试题39：数组中出现次数超过一半的数字的解法，利用partition函数。首先将数组按照基准进行调整，将数组分为两部分，左边都小于基准，右边都大于基准，这样如果基准恰好是第k个数字（从第0个数字开始计数），那么前k个数字就是要求的解。如果基准的下标小于k（下标从0开始），那么对基准右边应用partition函数；如果基准的下标大于k，则对基准左边应用partition函数。重复这一过程，直至基准恰好是第k个数字。代码如下：

class Solution {
public:
    
    void swap(vector<int>& v, int i, int j){
        int t = v[i];
        v[i] = v[j];
        v[j] = t;
    }
    
    int partition(vector<int>& v, int left, int right){
        int i = left;
        int j = right;
        int baseline = v[left];
        while(i<j){
            while(v[j]>=baseline && i<j)
                j--;
            while(v[i]<=baseline && i<j)
                i++;
            swap(v, i, j);
        }
        swap(v, left, i);
        return i;
    }
    
    vector<int> GetLeastNumbers_Solution(vector<int> input, int k) {
        vector<int> ans;
        if(input.empty() || k<=0 || k>input.size())
            return ans;
        
        int left = 0;
        int right = input.size()-1;
        int idx = partition(input, left, right);
        while(idx!=k){
            if(idx>k){
                right = idx-1;
                idx = partition(input, left, right);
            }
            else if(idx<k){
                left = idx+1;
                idx = partition(input, left, right);
            }
        }
        
        for(int i=0; i<k; i++)
            ans.push_back(input[i]);
        return ans;
    }
};

该算法的时间复杂度为O(n)，该算法不能保证输出的k个最小数字是排序的。

思路3

上面的2中做法都会修改输入的数组。如果不允许修改输入的数组，我们可以使用如下思路：使用一个容量为k容器保存前k个数，遍历输入的数组，如果容器未满，则将下一个元素放进容器中；如果容器已满，则找到容器中最大的元素：如果容器最大的元素大于下一个元素，则删除容器中最大的元素并将下一个元素放进容器，否则保留容器中的数据不变。当容器已满时，我们需要进行3个操作：查找最大值、删除和插入。可以使用最大堆来实现这3个操作，在最大堆中，需要O(1)的时间来得到最大值，需要O(logk)的时间来进行插入和删除。如果可以使用stl的话，可以使用stl中的multiset作为容器，multiset基于红黑树实现，其查找、插入、删除均需要O(logk)的时间。下面使用multiset实现这一算法：

class Solution {
public:
    vector<int> GetLeastNumbers_Solution(vector<int> input, int k) {
        vector<int> ans;
        if(input.empty() || k<=0 || k>input.size())
            return ans;
        
        multiset<int, greater<int>> container;
        for(int i=0; i<input.size(); i++){
            if(container.size()<k){
                container.insert(input[i]);
            }
            else{
                multiset<int, greater<int>>::iterator iter = container.begin();
                if(input[i]<*iter){
                    container.erase(iter);
                    container.insert(input[i]);
                }
            }
        }
        
        multiset<int, greater<int>>::iterator iter = container.begin();
        while(iter!=container.end()){
            ans.push_back(*iter);
            iter++;
        }
        return ans;
    }
};

该算法的时间复杂度为O(nlogk)，为3个算法中时间复杂度最低的。该算法无需对数据进行移动，比较适合海量数据。