Leetcode 347. 前 K 个高频元素

347. 前 K 个高频元素

题目：
给定一个非空的整数数组，返回其中出现频率前 k 高的元素。

示例 1:

输入: nums = [1,1,1,2,2,3], k = 2
输出: [1,2]
示例 2:

输入: nums = [1], k = 1
输出: [1]
说明：

你可以假设给定的 k 总是合理的，且 1 ≤ k ≤ 数组中不相同的元素的个数。
你的算法的时间复杂度必须优于 O(n log n) , n 是数组的大小。

---

分析：
TOP-K问题的几种解法：

• 各种排序，最好是快速排序O（nlgn）
• 快速选择算法 O(n)
• 堆排序：C++中大根堆的实现是优先队列，可以直接使用。O（lgn)
• BFPRT算法（最好最坏均是O(n)）

最简单的方法就是统计出现次数后调用STL的sort进行排序，其时间复杂度为O（nlgn），不符合题目要求，弃之。由于本题中不是简单的对原始数据求TOP-K，还需要HashMap对出现次数进行统计，在统计结果上求TOP-K，再映射回原始数据。所以虽然快速选择算法是一个很好的算法（ 时间复杂度：O(n) ），但是需要在几个数据结构中来回转换，会浪费一些空间，写起来也不方便，所以更好的做法是利用优先队列，时间复杂度主要体现在push()/pop()操作上：O（lgn）。在此选择两种解法，具体思路参见代码中的注释。

---

法一：优先队列（堆排序）

class Solution {
public:
    vector<int> topKFrequent(vector<int>& nums, int k) {
        unordered_map<int, int> map;
               // 统计每个元素出现次数，可保证不会出现重复元素
        for (auto n : nums) {
            map[n]++;
        }
        // 优先队列默认是一个大根堆，每次加入新的元素默认根据pair的first进行节点位置的调整（排序）
        priority_queue<pair<int, int>> q;
        for (auto m : map) {
            q.push(pair<int, int>(m.second, m.first));
        }
        vector<int> res(k, 0);
       // 每次取出的堆顶就是出现次数最多的元素，将其移除后堆又会自动调整
        for (int i = 0; i < k; i++) {
            res[i] = q.top().second;
            q.pop();
        }
        return res;
    }
};

法二：快速选择算法

class Solution {
public:
    // 需要交换的是整个pair对
    void swap(pair<int, int>& a, pair<int, int>& b) {
        pair<int, int> tmp = a;
        a = b;
        b = tmp;
    }

    // 下面两个加粗的函数，是快速选择通用的代码，求第K大和求前K大是一样的，因为partion的作用就是把比pivot大的都放在它左边，比它小的都放在它右边
    // 所以只要第K大求出来就意味着前面都是比它大的，即前K大求毕

    int partion(vector<pair<int, int> >& map, int left, int right) {
        int i = left;
        int j = right;
        // 按照出现次数进行排序，即.second
        int pivot = map[left].second;
        while (i < j) {
            while (i < j && map[j].second < pivot) j--;
            while (i < j && map[i].second >= pivot) i++;
            if (i < j) swap(map[i], map[j]);
        }
        swap(map[j], map[left]);
        return j;
    }
    void quickSelect(vector<pair<int, int> >& map, int left, int right, int k) {
        int p = partion(map, left, right);
        if (p == k - 1) return;
        if (p < k - 1) {
            quickSelect(map, p + 1, right, k);
        } else {
            quickSelect(map, left, p - 1, k);
        }
    }
    vector<int> topKFrequent(vector<int>& nums, int k) {
        unordered_map<int, int> map;
        for (auto n : nums) map[n]++;
        // 直接对map进行快速选择不方便，将数据转存到pair向量中
        vector<pair<int, int>> v;
        for (auto m : map) {
            v.push_back(pair<int, int>(m.first, m.second));
        }
        quickSelect(v, 0, v.size() - 1, k);
        for (int i = 0; i < k; i++) {
            nums[i] = v[i].first;
        }
        return vector<int>(nums.begin(), nums.begin() + k);
    }
};