leetcode347：前K个高频元素

题目链接

leetcode

题目描述

给你一个整数数组 nums 和一个整数 k ，请你返回其中出现频率前 k 高的元素。你可以按 任意顺序 返回答案。

C++ 代码

#include <iostream>
#include <vector>
#include <queue>
#include <unordered_map>
#include <algorithm>
using namespace std;

/*
    哈希表 + 堆
    自定义堆中保存的结构体（定义成员、定义比较规则）
    
    时间复杂度：O(nlogk)
        n 为数组的长度，k 为高频元素的个数。
    空间复杂度：O(n)
        哈希表的大小为 O(n)，堆的大小为 O(k)，共计为 O(n)。
*/
class Solution_0 {
public:
    vector<int> topKFrequent(vector<int>& nums, int k) {
        for (const int& num : nums) {
            num2cnt[num]++;
        }

        for (const auto& it : num2cnt) {
            heap.push({it.second, it.first});
            if (heap.size() > k) {
                heap.pop();
            }
        }

        vector<int> res;
        res.reserve(k);
        while (!heap.empty()) {
            Node node = heap.top(); heap.pop();
            res.push_back(node.num);
        }
        reverse(res.begin(), res.end()); // 升序排列

        return res;
    }

private:
    struct Node {
        int cnt;
        int num;
        bool operator<(const Node& t) const {
            if (cnt != t.cnt) return cnt > t.cnt;
            return num < t.num;
        }
    };

    priority_queue<Node> heap;       // 大根堆
    unordered_map<int, int> num2cnt; // 哈希表
};

/*
    思路同上
    堆中保存元素类型是 pair<int, int>，
    自定义 priority_queue 排序规则

    时间复杂度：O(nlogk)
    空间复杂度：O(n)
*/
class Solution {
public:
    vector<int> topKFrequent(vector<int>& nums, int k) {
        unordered_map<int, int> num2cnt;
        for (const int& num : nums) {
            num2cnt[num]++;
        }

        priority_queue<pair<int, int>, vector<pair<int, int>>, decltype(&cmp)> heap(cmp);
        for (const auto& it : num2cnt) {
            int num = it.first, cnt = it.second;
            heap.push({num, cnt});
            if (heap.size() > k) {
                heap.pop();
            }
        }

        // vector<int> res;
        // res.reserve(k);
        // while (!heap.empty()) {
        //     res.emplace_back(heap.top().first);
        //     heap.pop();
        // }
        // reverse(res.begin(), res.end()); // 升序排列
        vector<int> res(k);
        int i = k-1;
        while (!heap.empty()) {
            res[i] = heap.top().first;
            heap.pop();
            i--;
        }

        return res;
    }

private:
    static bool cmp(pair<int, int>& a, pair<int, int>& b) {
        return a.second > b.second;
    }
};

// 辅助函数：打印数组
void printArray(const vector<int>& nums) {
    cout << "[";
    for (size_t i = 0; i < nums.size(); i++) {
        cout << nums[i];
        if (i != nums.size() - 1) cout << ",";
    }
    cout << "]";
}

int main() {
    Solution solution;
    vector<vector<int>> nums_cases = {
        {1,1,1,2,2,3},
        {1}
    };
    vector<int> k_cases = {
        2, 1
    };

    for (int i = 0; i < nums_cases.size(); i++) {
        auto& nums = nums_cases[i];
        int k = k_cases[i];

        cout << "Input: nums = ";
        printArray(nums);
        cout << ", k = " << k << endl;
        vector<int> result = solution.topKFrequent(nums, k);
        cout << "Output: ";
        printArray(result);
        cout << endl;
    }

    return 0;
}