Problem statement

https://leetcode.com/problems/maximum-performance-of-a-team/

Solution

Let us look at pairs (efficiency, speed). Then what we can do first is to sort data by efficiency, starting from the biggest one and add element by element in our active set. By active set I mean the set of workers which efficiency is bigger than some given number. When we iterate through our data, active set can only be increasing: elements will be added, but never removed. At each moment of time we will ask the question: given smallest element in group of workers equal to say eff, what is the biggest performance we can get?

  1. We will keep our active set in heap: priority queue, because at each moment of time we need to choose k workers with the biggest efficiencies.
  2. sm is the sum of elements in our heap, ans is variable for our answer.
  3. We iterate over elements going from the highest efficiency and do the following steps: 3.1 Check if we have too much elements in our heap: if we have, remove the smallest one. 3.2 Add element to heap. At this stage we will have no more than k elements. 3.3 Update our sm: subtract popped element if needed and add new element 3.4 Update ans as maximum among ans and sm * eff, because by our construction eff is the smallest efficiency among our group.

Complexity

Let n be the length of S and E. Then we need O(n log n) to sort our data, then we need O(n log k) to work with heap updates. So overall time complexity is O(n log n + n log k). Space complexity is O(n + k).

Code

class Solution:
    def maxPerformance(self, n, S, E, k):
        sm, ans, heap = 0, 0, []
        
        for eff, speed in sorted(zip(E, S))[::-1]:
            if len(heap) > k - 1: sm -= heappop(heap)
            heappush(heap, speed)
            sm += speed
            ans = max(ans, sm*eff)
        
        return ans % (10**9 + 7)