【快速选择算法】解决TopK问题中前K小的数字问题

1.前言

在一个数组中找到这个数组前K小的数字有三种方式：

• 排序 O(N*logN)
• 堆排序：建立一个k个大小的大堆(如果是找前K大的数字的话用小堆) O(N*logK)
• 快速选择算法：原地交换数字，使得该数组前k个数字就是该数组前K小的数字 一般情况接近O(N)，最差情况O(N^2)

2.题目简介

题目链接：LINK

3.求解思路

4.示例代码

class Solution {
public:
    vector<int> inventoryManagement(vector<int>& nums, int k) 
    {
        srand(time(nullptr));

        qsort(nums, 0, nums.size() - 1, k);
        return {nums.begin(), nums.begin() + k};
    }

    void qsort(vector<int>& nums, int begin, int end, int k)
    {
        if(begin >= end) return;

        //数组分三块
        int key = GetRandom(nums, begin, end);
        int left = begin - 1, right = end + 1, i = begin;
        while(i < right)
        {
            if(nums[i] < key)
            {
                swap(nums[i++], nums[++left]);
            }
            else if(nums[i] == key)
            {
                i++;
            }
            else //nums[i] > key
            {
                swap(nums[i], nums[--right]);
            }
        }

        //分情况讨论
        int a = left - begin + 1;
        int b = right - 1 - (left + 1) + 1;
        int c = end - begin + 1 - a - b;
        if(a >= k) qsort(nums, begin, left, k);
        else if(a + b >= k) return;
        else qsort(nums, right, end, k - a - b);
    }
    int GetRandom(vector<int>& nums, int begin, int end)
    {
        return nums[rand() % (end - begin + 1) + begin];
    }
};

EOF