LeetCode-in-All

347. Top K Frequent Elements

Medium

Given an integer array nums and an integer k, return the k most frequent elements. You may return the answer in any order.

Example 1:

Input: nums = [1,1,1,2,2,3], k = 2

Output: [1,2]

Example 2:

Input: nums = [1], k = 1

Output: [1]

Constraints:

Follow up: Your algorithm’s time complexity must be better than O(n log n), where n is the array’s size.

Solution

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

class Solution {
public:
    vector<int> topKFrequent(vector<int>& nums, int k) {
        sort(nums.begin(), nums.end());

        // Min heap of <number, frequency>
        priority_queue<pair<int, int>, vector<pair<int, int>>, greater<pair<int, int>>> minHeap;

        // Filter with min heap
        int j = 0;
        for (int i = 0; i <= nums.size(); ++i) {
            if (i == nums.size() || nums[i] != nums[j]) {
                minHeap.push({i - j, nums[j]});
                if (minHeap.size() > k) {
                    minHeap.pop();
                }
                j = i;
            }
        }

        // Convert to int array
        vector<int> result(k);
        for (int i = k - 1; i >= 0; --i) {
            result[i] = minHeap.top().second;
            minHeap.pop();
        }

        return result;
    }
};