剑指40：最小的k个数

传送门

nowcoder
leetcode

题目描述

给定一个长度为 n 的可能有重复值的数组，找出其中不去重的最小的 k 个数。
例如：数组元素是4,5,1,6,2,7,3,8这8个数字，则最小的4个数字是1,2,3,4（任意顺序）。
数据范围：0 <= k, n <= 10000，数组中每个元素的大小 0 <= val <= 1000
要求：空间复杂度 O(n)，时间复杂度 O(nlogk)

C++ 代码 - nowcoder

/*
    优先队列，小顶堆

    时间复杂度：O(n*logk)
*/
class Solution {
public:
    vector<int> GetLeastNumbers_Solution(vector<int>& input, int k) {
        if (k > input.size()) return {};

        priority_queue<int, vector<int>, greater<int>> pq; // 升序队列，小顶堆
        for (auto x : input) {
            pq.push(x);
        }
        vector<int> res;
        while(k --) {
            res.push_back(pq.top());
            pq.pop();
        }

        return res;
    }
};

C++ 代码 - leetcode

/*
    优先队列，大顶堆（需要添加控制逻辑）
    priority_queue 默认是大顶堆 

    时间复杂度：O(n*logk)
*/
class Solution {
public:
    vector<int> inventoryManagement(vector<int>& arr, int k) {
        if (arr.empty() || k == 0) return vector<int>();
        if (arr.size() < k) return vector<int>();

        // priority_queue<int> pq;
        priority_queue<int, vector<int>, less<int>> pq; // 大顶堆
        // 前 k 个数构成最大堆
        for (int x : arr) {
            pq.push(x);
            // 堆中元素多于 k 个时，删除堆顶元素
            if (pq.size() > k) {
                pq.pop();
            }
        }
        // 堆中元素保存到数组
        vector<int> res(k);
        int i = 0;
        while (!pq.empty()) {
            res[i] = pq.top(); pq.pop();
            i ++;
        }

        return res;
    }
};

/*
	快速选择
	
	时间复杂度：O(n)
*/
class Solution {
public:
    vector<int> inventoryManagement(vector<int>& arr, int k) {
        vector<int> res(k);

        // 随机打乱数组
        shuffle(arr);
        int l = 0, r = arr.size() - 1;
        // 寻找第 k 大的元素
        while (l <= r) {
            // arr[l..r] 中选择一个分界点
            int p = partition(arr, l, r);
            if (p < k - 1) {
                // 第 k 大的元素在 arr[p+1..r]
                l = p + 1;
            } else if (p > k - 1) {
                // 第 k 大的元素在 arr[l..p-1]
                r = p - 1;
            } else {
                // arr[p] 就是第 k 大元素，左边的元素都比它小，arr[0..k] 就是答案
                for (int i = 0; i < k; i ++) {
                    res[i] = arr[i];
                }
                return res;
            }
        }
        
        return res;
    }
    
    // 洗牌算法，将输入的数组随机打乱
    void shuffle(vector<int>& nums) {
        srand((unsigned)time(NULL));
        int n = nums.size();
        for (int i = 0; i < n; i ++) {
            // 生成 [i, n - 1] 的随机数
            int r = i + rand() % (n - i);
            swap(nums[i], nums[r]);
        }
    }

    // 对 nums[l...r] 进行切分
    int partition(vector<int>& nums, int l, int r) {
        int pivot = nums[l];
        // i, j 是开区间，定义为：[l, i) <= pivot；(j, r] > pivot
        // 维护这个区间的定义
        int i = l + 1, j = r;
        while (i <= j) {
            while (i < r && nums[i] <= pivot) {
                i ++; // while 结束时恰好 nums[i] > pivot
            }
            while (j > l && nums[j] > pivot) {
                j --; // while 结束时恰好 nums[j] <= pivot
            }
            if (i >= j) {
                break;
            }
            // 维护区间定义
            swap(nums[i], nums[j]);
        }
        // 将 pivot 放到合适的位置，即 pivot 左边元素较小，右边元素较大
        swap(nums[l], nums[j]);

        return j;
    }

    // 原地交换数组中的两个元素
    void swap(int& a, int& b) {
        int temp = a;
        a = b;
        b = temp;
    }
};