侧边栏壁纸
博主头像
GabrielxD

列車は必ず次の駅へ。では舞台は?私たちは?

  • 累计撰写 471 篇文章
  • 累计创建 108 个标签
  • 累计收到 9 条评论

目 录CONTENT

文章目录

【优先队列, 排序, BST】最小的k个数

GabrielxD
2022-09-17 / 0 评论 / 0 点赞 / 24 阅读 / 810 字 / 正在检测是否收录...

题目

剑指 Offer 40. 最小的k个数


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

示例 1:

输入:arr = [3,2,1], k = 2
输出:[1,2] 或者 [2,1]

示例 2:

输入:arr = [0,1,2,1], k = 1
输出:[0]

限制:

  • 0 <= k <= arr.length <= 10000
  • 0 <= arr[i] <= 10000

解题

方法一:大根堆

思路

本题是经典的 Top K 问题(前 K 小),维护一个容量为 k 的大根堆(maxHeap)。
先用数组中前 k 个元素把优先队列填满,接着遍历数组,如果当前元素小于堆顶,就把堆顶元素出队,然后放入当前元素,否则什么都不干。
遍历完成后把优先队列中所有元素取出放入数组中返回即可。

代码

class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        if (k == 0) return new int[0];
        Queue<Integer> maxHeap = new PriorityQueue<>(Comparator.reverseOrder());
        for (int i = 0; i < k; ++i) maxHeap.offer(arr[i]);
        for (int i = k; i < arr.length; ++i) {
            if (arr[i] < maxHeap.peek()) {
                maxHeap.poll();
                maxHeap.offer(arr[i]);
            }
        }
        int[] ans = new int[k];
        for (int i = 0; i < k; ++i) ans[i] = maxHeap.poll();
        return ans;
    }
}
class Solution {
public:
    vector<int> getLeastNumbers(vector<int>& arr, int k) {
        if (!k) return { };
        priority_queue<int> max_heap;
        for (int i = 0; i < k; ++i) max_heap.push(arr[i]);
        for (int i = k; i < arr.size(); ++i) {
            if (arr[i] < max_heap.top()) {
                max_heap.pop();
                max_heap.push(arr[i]);
            }
        }
        vector<int> ans(k);
        for (int i = 0; i < k; ++i) ans[i] = max_heap.top(), max_heap.pop();
        return ans;
    }
};

方法二:排序

思路

直接把数组升序排序,然后返回前 k 个元素即可。

代码

class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        Arrays.sort(arr);
        return Arrays.copyOf(arr, k);
    }
}
class Solution {
public:
    vector<int> getLeastNumbers(vector<int>& arr, int k) {
        sort(arr.begin(), arr.end());
        return {arr.begin(), arr.begin() + k};
    }
};

本题数据范围比较小,直接上计数排序。

class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        int[] counts = new int[10001];
        for (int num : arr) ++counts[num];
        int[] topK = new int[k];
        for (int i = 0; i <= 10000 && k > 0; ++i) {
            while (counts[i]-- > 0 && k > 0) topK[--k] = i;
        }
        return topK;
    }
}
class Solution {
public:
    vector<int> getLeastNumbers(vector<int>& arr, int k) {
        int counts[10001];
        memset(counts, 0, sizeof(counts));
        for (int& num : arr) ++counts[num];
        vector<int> top_k;
        for (int i = 0; i <= 10000 && k > 0; ++i) {
            while (counts[i]-- && k) {
                top_k.push_back(i);
                --k;
            }
        }
        return top_k;
    }
};

方法三:二叉搜索树(BST)

思路

思路与上面差不多,也是利用二叉搜索树进行排序。

C++ 直接用 multiset,Java 里没有可重复的有序集合,使用 TreeMap 代替。

代码

class Solution {
public:
    vector<int> getLeastNumbers(vector<int>& arr, int k) {
        multiset<int> bst(arr.begin(), arr.end());
        auto pos_k = bst.begin();
        advance(pos_k, k);
        return vector<int>(bst.begin(), pos_k);
    }
};
class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        Map<Integer, Integer> bst = new TreeMap<Integer, Integer>();
        for (int num : arr) bst.put(num, bst.getOrDefault(num, 0) + 1);
        int[] topK = new int[k];
        outer: for (var entry : bst.entrySet()) {
            int num = entry.getKey(), cnt = entry.getValue();
            while (cnt-- > 0) {
                if (--k < 0) break outer;
                topK[k] = num;
            }
        }
        return topK;
    }
}
0

评论区