Leetcode 347 - Top K Frequent Elements

Link to problem

Problem

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

Solution

I have two slightly different solutions to the problem. But the preparation works are the same: First count the numbers with a map, then

• Array (sort): Dump all the contents into an array, then sort it by frequency.
• Min Heap: Put the pairs in a heap. The heap could either be a max heap, which you get k numbers from top. Or you use a min heap, but you pop the top element if the heap size is larger than k. The method using min heap is more efficient since the time complexity of pushing / popping elements could be lower. Moreover, memory consumption is also better as there are at most, k+1 elements stored in the heap.

Implementation

A. Sorting array

vector<int> topKFrequent(vector<int>& nums, int k) {
    unordered_map<int, int> bin;
    for (int n: nums) bin[n]++;

    vector<pair<int, int>> dump(bin.begin(), bin.end());
    sort(dump.begin(), dump.end(), [](auto lhs, auto rhs)->bool {
        return lhs.second > rhs.second;
    });

    vector<int> result;
    for (int i = 0; i < k; i++) result.push_back(dump[i].first);

    return result;
}

B. Heap (min-heap)

Min heap in C++ could be done using priority_queue.

Constructing priority_queue

I found a solution that constructs the priority_queue with a lambda function comparator in a pleasant way. What I usually do is:

// binIt is the iterator of the unordered_map that stores the number and its frequency
auto comp = [](binIt lhs, binIt rhs) -> bool { return lhs->second > rhs->second; };
priority_queue(binIt, vector<binIt>, decltype(comp)) heap(comp, vector<binIt>());

Since I thought I couldn’t get the type of a lambda function, I will need decltype to slap it in priority_queue’s template. However, there is a wrapper called std::function that does the job. Therefore, declaration of the priority_queue becomes:

priority_queue<binIt, vector<binIt>, function<bool(binIt, binIt)>> heap
    ([](binIt lhs, binIt rhs) -> bool { return lhs->second > rhs->second; },
        vector<binIt>());

Now the whole comparator could be fitted into the parameter of priority_queue.

The full implementation using min heap is as the following:

vector<int> topKFrequent(vector<int>& nums, int k) {
    using binIt = unordered_map<int, int>::iterator;

    unordered_map<int, int> bin;
    for (int n: nums) bin[n]++;

    // min-heap
    priority_queue<binIt, vector<binIt>, function<bool(binIt, binIt)>> heap
        ([](binIt lhs, binIt rhs) -> bool { return lhs->second > rhs->second; },
            vector<binIt>());

    for (binIt it = bin.begin(); it != bin.end(); it++) {
        heap.push(it);
        if (heap.size() > k) heap.pop();
    }

    vector<int> result;
    while (!heap.empty()) {
        result.push_back(heap.top()->first);
        heap.pop();
    }

    return result;
}