剑指offer-获取数组前K小的数字

/*******************************************************************
Copyright(c) 2016, Tyrone Li
All rights reserved.
*******************************************************************/
// 作者:TyroneLi
//

/*
Q:
    最小的K个数：
        输入n个整数，找出其中最小的k个数。
S:
    1. 借鉴快排中的每次排序的分治函数，找出前k个最小的数，那么当进行的某次快排到达(k-1)的位置时候，
       前k个数就在此时的位置及其前面的(k-1)个数字。
    2. 由于方法1会改变原来的数组的数字位置，这里采用一个大小k的辅助存储空间，当这个空间没有存满k个数
       字的时候，就从原先的数组从前往后遍历并插入，存满k个数字之后，遍历得到的当前数字与存储空间的最大
       数字比较，如果比存储空间最大的数字还小，就替换；为了满足这样的操作而且时间复杂度不高，选择stl库
       中的set或者multiset，因为其底层是采用红黑树来实现的。

*/

#include <iostream>
#include <set>
#include <cstdio>
#include <cstdlib>

typedef std::multiset<int, std::greater<int> >            inSet;
typedef std::multiset<int, std::greater<int> >::iterator  setIter;

int partition(int*numbers, int low, int high)
{
    if(low < high)
    {
        int pivort = numbers[low];
        while(low < high)
        {
            while(low < high && numbers[high] >= pivort) high--;
            numbers[low] = numbers[high];
            while(low < high && numbers[low] <= pivort) low++;
            numbers[high] = numbers[low];
        }
        numbers[low] = pivort;
    }
    return low;
}

void getLeastKNums_partition(int*numbers, int*output, int length, int k)
{
    if(numbers == nullptr || output == nullptr || k > length || k <= 0 || length <= 0)
        return;
    
    int start = 0;
    int end = length - 1;
    int index = partition(numbers, start, end);
    while(index != (k - 1))
    {
        if(index > (k - 1))
        {
            end = index - 1;
            index = partition(numbers, start, end);
        }else{
            start = index + 1;
            index = partition(numbers, start, end);
        }
    }
    for(int i = 0; i < k; ++i)
        output[i] = numbers[i];
}

void getLeastKNums_set(int*numbers, inSet&leastNums, int length, int k)
{
    if(numbers == nullptr || k < 1 || k > length || length <= 0)
        return;

    for(int i = 0; i < length; ++i)
    {
        if(leastNums.size() < k)
        {
            leastNums.insert(numbers[i]);
        }else{
            setIter inSetLeastOne = leastNums.begin();
            if((*inSetLeastOne) > numbers[i])
            {
                leastNums.erase(inSetLeastOne);
                leastNums.insert(numbers[i]);
            }
        }
    }
}

void test_1()
{
    std::cout << "Test 1" << std::endl;
    int numbers[] = {4,2,3,7,1,6,5,8,9,10};
    std::cout << "input arr:" << std::endl;
    for(auto e:numbers)
        std::cout << e << " ";
    std::cout << std::endl;
    int*output = new int[4];
    getLeastKNums_partition(numbers, output, 10, 4);
    for(int i = 0; i < 4; ++i)
        std::cout << output[i] << " ";
    std::cout << std::endl;

    int numbers_set[] = {4,2,3,7,1,6,5,8,9,10};
    inSet output_set;
    getLeastKNums_set(numbers_set, output_set, 10, 4);
    setIter iter = output_set.begin();
    for(; iter != output_set.end(); ++iter)
        std::cout << *iter << " ";
    std::cout << std::endl;
}

void test_2()
{
    std::cout << "Test 2" << std::endl;
    int numbers[] = {4,2,4,7,1,6,5,8,9,10};
    std::cout << "input arr:" << std::endl;
    for(auto e:numbers)
        std::cout << e << " ";
    std::cout << std::endl;
    
    int*output = new int[4];
    getLeastKNums_partition(numbers, output, 10, 4);
    for(int i = 0; i < 4; ++i)
        std::cout << output[i] << " ";
    std::cout << std::endl;

    int numbers_set[] = {4,2,4,7,1,6,5,8,9,10};
    inSet output_set;
    getLeastKNums_set(numbers_set, output_set, 10, 4);
    setIter iter = output_set.begin();
    for(; iter != output_set.end(); ++iter)
        std::cout << *iter << " ";
    std::cout << std::endl;
}

void test_3()
{
    std::cout << "Test 3" << std::endl;
    int numbers[] = {4,2,3,7,1,6,6,8,9,10};
    std::cout << "input arr:" << std::endl;
    for(auto e:numbers)
        std::cout << e << " ";
    std::cout << std::endl;
    
    int*output = new int[4];
    getLeastKNums_partition(numbers, output, 10, 4);
    for(int i = 0; i < 4; ++i)
        std::cout << output[i] << " ";
    std::cout << std::endl;

    int numbers_set[] = {4,2,3,7,1,6,6,8,9,10};
    inSet output_set;
    getLeastKNums_set(numbers_set, output_set, 10, 4);
    setIter iter = output_set.begin();
    for(; iter != output_set.end(); ++iter)
        std::cout << *iter << " ";
    std::cout << std::endl;
}

void test_4()
{
    std::cout << "Test 4" << std::endl;
    int numbers[] = {4,2,3,7,1,6,5,8,9,10};
    std::cout << "input arr:" << std::endl;
    for(auto e:numbers)
        std::cout << e << " ";
    std::cout << std::endl;
    
    int*output = new int[4];
    getLeastKNums_partition(numbers, output, 10, 0);
    for(int i = 0; i < 4; ++i)
        std::cout << output[i] << " ";
    std::cout << std::endl;

    int numbers_set[] = {4,2,3,7,1,6,5,8,9,10};
    inSet output_set;
    getLeastKNums_set(numbers_set, output_set, 10, 0);
    setIter iter = output_set.begin();
    for(; iter != output_set.end(); ++iter)
        std::cout << *iter << " ";
    std::cout << std::endl;
}

void test_5()
{
    std::cout << "Test 5" << std::endl;
    int numbers[] = {4,2,3,7,1,6,5,8,9,10};
    std::cout << "input arr:" << std::endl;
    for(auto e:numbers)
        std::cout << e << " ";
    std::cout << std::endl;
    
    int*output = new int[4];
    getLeastKNums_partition(numbers, output, 10, 1);
    for(int i = 0; i < 4; ++i)
        std::cout << output[i] << " ";
    std::cout << std::endl;

    int numbers_set[] = {4,2,3,7,1,6,5,8,9,10};
    inSet output_set;
    getLeastKNums_set(numbers_set, output_set, 10, 1);
    setIter iter = output_set.begin();
    for(; iter != output_set.end(); ++iter)
        std::cout << *iter << " ";
    std::cout << std::endl;
}

void test_6()
{
    std::cout << "Test 6" << std::endl;
    int numbers[] = {4,2,3,7,1,6,5,8,9,10};
    std::cout << "input arr:" << std::endl;
    for(auto e:numbers)
        std::cout << e << " ";
    std::cout << std::endl;
    
    int*output = new int[4];
    getLeastKNums_partition(numbers, output, 10, 11);
    for(int i = 0; i < 4; ++i)
        std::cout << output[i] << " ";
    std::cout << std::endl;

    int numbers_set[] = {4,2,3,7,1,6,5,8,9,10};
    inSet output_set;
    getLeastKNums_set(numbers_set, output_set, 10, 11);
    setIter iter = output_set.begin();
    for(; iter != output_set.end(); ++iter)
        std::cout << *iter << " ";
    std::cout << std::endl;
}

void test_getLeastKNums()
{
    test_1();
    test_2();
    test_3();
    test_4();
    test_5();
    test_6();
}

int main(int argc, char**argv)
{

    test_getLeastKNums();

    return 0;
}