Leetcode 1792 - Maximum Average Pass Ratio

Link to Problem

Problem

There is a school that has classes of students and each class will be having a final exam. You are given a 2D integer array classes, where classes [i] = [passi, totali]. You know beforehand that in the ith class, there are totali total students, but only passi number of students will pass the exam.

You are also given an integer extraStudents. There are another extraStudents brilliant students that are guaranteed to pass the exam of any class they are assigned to. You want to assign each of the extraStudents students to a class in a way that maximizes the average pass ratio across all the classes.

The pass ratio of a class is equal to the number of students of the class that will pass the exam divided by the total number of students of the class. The average pass ratio is the sum of pass ratios of all the classes divided by the number of the classes.

Return the maximum possible average pass ratio after assigning the extraStudents students.

Approach

This problem is quite straight forward, to maximize the average pass ratio, we need to get the class that could improve the most when adding one to the pass ratio’s numerator and denominator. Therefore, a heap, or std::priority_queue could help us achieve this. The rest could be done in a greedy manner.

Solution

1. Form a heap sorted by the maximum possible improvement of class.
2. For number of extraStudents, pop the class from the top of the heap and add one to its nominator and denominator. Push back.
3. Accumulate all class’ pass ratio and return the average.

Implementation

Special Note: If using std::vector<int, int> to store the classes, it will result in TLE for some reason. The very naive guess for this is because std::pair is a more simpler implementation in STL compared to std::vector.

class Solution {
public:
    double maxAverageRatio(vector<vector<int>>& classes, int extraStudents) {
        auto compare = [](const pair<int, int>& a, const pair<int, int>& b) {
            double diffA = ((double)(a.first+1)/(a.second+1)) - ((double)a.first/a.second);
            double diffB = ((double)(b.first+1)/(b.second+1)) - ((double)b.first/b.second);
            return diffA < diffB;
        };

        priority_queue<pair<int, int>, vector<pair<int, int>>, decltype(compare)> heap(compare);

        for (auto c: classes) {
            heap.push({c[0], c[1]});
        }

        for (int i = 0; i < extraStudents; i++) {
            auto temp = heap.top();
            heap.pop();
            temp.first++; temp.second++;
            heap.push(temp);
        }

        double sum = 0;
        while (!heap.empty()) {
            sum += ((double)heap.top().first / (double)heap.top().second);
            heap.pop();
        }

        return sum / classes.size();
    }
};