最小的k个数

题目描述

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

思路一：基于partition函数的算法，时间复杂度可达到O(n)

class Solution
{
    vector<int> GetLeastNumbers_Solution( vector<int> input, int k )
    {
        if ( k <= 0 )
            return vector<int>();
        if ( input.size() < k )
            return vector<int>();
        int start = 0, end = input.size()-1;
        int kIndex = k - 1;
        int index;
        vector<int> output;
        while ( kIndex != ( index = partition( input, start, end ) ) ) 
        {
            if ( kIndex > index )
                start = index+1;
            else 
                end = index-1;
        }
        for ( int i = 0; i < k; i++ )
        {
            output.push_back( input[i] );
        }
        return output;

    }
    //[start, end]
    int partition( vector<int> &input, int start, int end )
    {
        int sentry = start;
        start++;
        while ( start <= end )
        {
            if ( input[start] <= input[sentry] )
            {
                start++;
                continue;
            }
            if ( input[end] > input[sentry] )
            {
                end--;
                continue;
            }
            swap( &input[start], &input[end] );
            start++;
            end--;
        }
        swap( &input[end], &input[sentry] );
        return end;
    }
    void swap( int *a, int *b )
    {
        int temp = *a;
        *a = *b;
        *b = temp;
    }
}

思路二：借助最大堆时间复杂度可达到O(nlogk)的算法：

class Solution {
public:
    vector<int> GetLeastNumbers_Solution( vector<int> input, int k )
    {
        intSet leastNumbers;
        if ( k < 1 || k > input.size() )
            return vector<int>();
        vector<int>::const_iterator iter = input.begin();
        for ( ; iter != input.end(); ++iter )
        {
            if ( leastNumbers.size() < k )
                leastNumbers.insert( *iter );
            else 
            {
                setIterator iterGreatest = leastNumbers.begin();

                if ( *iter < *(leastNumbers.begin()) )
                {
                    leastNumbers.erase( iterGreatest );
                    leastNumbers.insert( *iter );
                }
            }
        }
        return vector<int>( leastNumbers.begin(), leastNumbers.end() );
    }
private:
    typedef multiset<int, greater<int> >    intSet;
    typedef multiset<int, greater<int> >::iterator setIterator;

};