Priority Queue

Account Login - LeetCode

class Solution:
    def maxPerformance(self, n: int, speed: List[int], efficiency: List[int], k: int) -> int:
        if n==1:return speed[0]*efficiency[0]
        if k==1:return max(a*b for a,b in zip(speed, efficiency))
        spe_eff = [(a,b) for a,b in zip(speed, efficiency)]
        spe_eff.sort(key=lambda x:x[1])
        pq, ksum = [], 0
        heapq.heapify(pq)
        # at most K. We initialize the res as picking up the last one in the spe_eff list
        res = spe_eff[-1][0]*spe_eff[-1][1]
        for i in range(n-1, 0, -1):
            spe, eff = spe_eff[i]
            # if the priority queue's length is less than k-1, we can add more elements into the queue. 
            if len(pq)<k-1:
                ksum += spe
                heapq.heappush(pq,spe)
            # top K-1 element operation
            # if we reach the k-1 capacity, if new explored element is larger than the smallest element in the priority queue, we pop the smaller one and push the current one
            elif spe>pq[0]:
                ksum += spe - pq[0]
                heapq.heappushpop(pq,spe)
            this_res = (ksum+spe_eff[i-1][0])*spe_eff[i-1][1]
            res = max(res, this_res)
        return res%(10**9 + 7)

Above is my ugly code and below is the elegent code from lee215

Time complexity O(nlogk)

Space complexity O(k)

class Solution:
    def maxPerformance(self, n: int, speed: List[int], efficiency: List[int], k: int) -> int:
        h = []
        res = ksum = 0
        for e, s in sorted(zip(efficiency, speed), reverse=1):
            heapq.heappush(h, s)
            ksum += s
            if len(h) > k:
                ksum -= heapq.heappop(h)
            res = max(res, ksum * e)
        return res % (10**9 + 7)