Problem Statement in English

You’re given a list of classes, where $class[i] = [p_i, t_i]$.

Each class has a number of students who passed ($p_i$) and a total number of students ($t_i$).

Your task is to maximize the average pass ratio after distributing a certain number of extra students among the classes.

Each extra student is guaranteed to pass any class that they’re placed in.

Approach

The naive approach is to simulate the process of adding extra students one by one to the class with the highest potential gain in pass ratio.

This involves repeatedly finding the class that would benefit the most from an additional student, updating its pass and total counts, and recalculating the average pass ratio.

But this can be optimised by precalculating the gain that any extra student would provide to each class, and saving this to a maxheap, so that every time we have an extra student that we can add, we just pop the class with the highest gain from the heap, add the student, and push the updated class back onto the heap.

Solution in Python

  • Naive

class Solution:
    def maxAverageRatio(self, classes: List[List[int]], extraStudents: int) -> float:
        while extraStudents:
            gain = {}
            bestGain = bestIndex = 0

            for i in range(len(classes)):
                p, t = classes[i]

                gain[i] = (t - p)/(t * (t + 1))
                if bestGain < gain[i]:
                    bestIndex = i
                    bestGain = gain[i]

            extraStudents -= 1
            classes[bestIndex][0] += 1
            classes[bestIndex][1] += 1

        ratios = 0
        for p, t in classes:
            ratios += p/t

        return ratios/len(classes)
  • Optimised

class Solution:
    def maxAverageRatio(self, classes: List[List[int]], extraStudents: int) -> float:
        maxheap = []
        for p, t in classes:
            heappush(maxheap, (-(t - p)/(t * (t + 1)), p, t))

        heapify(maxheap)

        for _ in range(extraStudents):
            _, p, t = maxheap[0]
            p += 1
            t += 1
            heapreplace(maxheap, (-(t - p)/(t * (t + 1)), p, t))
            pass

        ratios = 0
        for _, p, t in maxheap:
            ratios += p/t

        return ratios/len(classes)

Complexity

  • Time: $O(k * \log(n))$
    Here, $k$ is the number of extra students and $n$ is the number of classes. Each time we add an extra student, we need to update the max heap, which takes $O(\log(n))$ time.

  • Space: $O(n)$
    We are using a max heap to store the classes, which can take up to $O(n)$ space.

Mistakes I Made

I couldn’t come up with the gain approach on my own. Neither could I come up with the idea of precomputing and pushing the updated gain value of every class in one shot.

And we are done.