Heap and priority queue in Python

Posted on Edited on In , Views:

Intro

In Python, “heapq” module is available to use. Whenever elements are pushed or popped, heap structure is maintained. The heap[0] element also returns the smallest element each time.

[Leetcode 857] Minimum Cost to Hire K Workers

There are n workers. You are given two integer arrays quality and wage where quality[i] is the quality of the ith worker and wage[i] is the minimum wage expectation for the ith worker.

We want to hire exactly k workers to form a paid group. To hire a group of k workers, we must pay them according to the following rules:

  • Every worker in the paid group should be paid in the ratio of their quality compared to other workers in the paid group.
  • Every worker in the paid group must be paid at least their minimum wage expectation.

Given the integer k, return the least amount of money needed to form a paid group satisfying the above conditions. Answers within 10-5 of the actual answer will be accepted.

Example 1:

1
2
3
Input: quality = [10,20,5], wage = [70,50,30], k = 2
Output: 105.00000
Explanation: We pay 70 to 0th worker and 35 to 2nd worker.

Example 2:

1
2
3
Input: quality = [3,1,10,10,1], wage = [4,8,2,2,7], k = 3
Output: 30.66667
Explanation: We pay 4 to 0th worker, 13.33333 to 2nd and 3rd workers separately.

Constraints:

  • n == quality.length == wage.length
  • 1 <= k <= n <= 104
  • 1 <= quality[i], wage[i] <= 104

Solution:

1
2
3
4
5
6
7
8
9
10
11
12
class Solution:
    def mincostToHireWorkers(self, quality: List[int], wage: List[int], k: int) -> float:
        res, qualitySum, maxHeap = float('inf'), 0, []

        for ration, q in sorted([(w / q, q) for w, q in zip(wage, quality)]):
            qualitySum += q
            heappush(maxHeap, -q)
            if len(maxHeap) == k:
                res = min(res, qualitySum * ration)
                qualitySum -= -heappop(maxHeap)
        
        return res

[Leetcode 1383] Maximum Performance of a Team

You are given two integers n and k and two integer arrays speed and efficiency both of length n. There are n engineers numbered from 1 to n. speed[i] and efficiency[i] represent the speed and efficiency of the ith engineer respectively.

Choose at most k different engineers out of the n engineers to form a team with the maximum performance.

The performance of a team is the sum of its engineers’ speeds multiplied by the minimum efficiency among its engineers.

Return the maximum performance of this team. Since the answer can be a huge number, return it modulo 109 + 7.

Example 1:

1
2
3
4
Input: n = 6, speed = [2,10,3,1,5,8], efficiency = [5,4,3,9,7,2], k = 2
Output: 60
Explanation: 
We have the maximum performance of the team by selecting engineer 2 (with speed=10 and efficiency=4) and engineer 5 (with speed=5 and efficiency=7). That is, performance = (10 + 5) * min(4, 7) = 60.

Example 2:

1
2
3
4
Input: n = 6, speed = [2,10,3,1,5,8], efficiency = [5,4,3,9,7,2], k = 3
Output: 68
Explanation:
This is the same example as the first but k = 3. We can select engineer 1, engineer 2 and engineer 5 to get the maximum performance of the team. That is, performance = (2 + 10 + 5) * min(5, 4, 7) = 68.

Example 3:

1
2
Input: n = 6, speed = [2,10,3,1,5,8], efficiency = [5,4,3,9,7,2], k = 4
Output: 72

Constraints:

  • 1 <= k <= n <= 105
  • speed.length == n
  • efficiency.length == n
  • 1 <= speed[i] <= 105
  • 1 <= efficiency[i] <= 108

Solution:

1
2
3
4
5
6
7
8
9
10
11
12
class Solution:
    def maxPerformance(self, n: int, speed: List[int], efficiency: List[int], k: int) -> int:
        res, speedSum, minHeap = 0, 0, []
        
        for ef, sp in sorted(zip(efficiency, speed), reverse=True):
            speedSum += sp
            heappush(minHeap, sp)
            res = max(res, speedSum * ef)
            if len(minHeap) == k:
                speedSum -= heappop(minHeap)
        
        return res % (10**9 + 7)

[Leetcode 2462] Total Cost to Hire K Workers

You are given a 0-indexed integer array costs where costs[i] is the cost of hiring the ith worker.

You are also given two integers k and candidates. We want to hire exactly k workers according to the following rules:

  • You will run k sessions and hire exactly one worker in each session.
  • In each hiring session, choose the worker with the lowest cost from either the first candidates(instead, read candidates as n to make it easier understand) workers or the last candidates workers. Break the tie by the smallest index.
    • For example, if costs = [3,2,7,7,1,2] and candidates = 2, then in the first hiring session, we will choose the 4th worker because they have the lowest cost [3,2,7,7,1,2].
    • In the second hiring session, we will choose 1st worker because they have the same lowest cost as 4th worker but they have the smallest index [3,2,7,7,2]. Please note that the indexing may be changed in the process.
  • If there are fewer than candidates workers remaining, choose the worker with the lowest cost among them. Break the tie by the smallest index.
  • A worker can only be chosen once.

Return the total cost to hire exactly k workers.

Example 1:

1
2
3
4
5
6
7
Input: costs = [17,12,10,2,7,2,11,20,8], k = 3, candidates = 4
Output: 11
Explanation: We hire 3 workers in total. The total cost is initially 0.
- In the first hiring round we choose the worker from [17,12,10,2,7,2,11,20,8]. The lowest cost is 2, and we break the tie by the smallest index, which is 3. The total cost = 0 + 2 = 2.
- In the second hiring round we choose the worker from [17,12,10,7,2,11,20,8]. The lowest cost is 2 (index 4). The total cost = 2 + 2 = 4.
- In the third hiring round we choose the worker from [17,12,10,7,11,20,8]. The lowest cost is 7 (index 3). The total cost = 4 + 7 = 11. Notice that the worker with index 3 was common in the first and last four workers.
The total hiring cost is 11.

Example 2:

1
2
3
4
5
6
7
Input: costs = [1,2,4,1], k = 3, candidates = 3
Output: 4
Explanation: We hire 3 workers in total. The total cost is initially 0.
- In the first hiring round we choose the worker from [1,2,4,1]. The lowest cost is 1, and we break the tie by the smallest index, which is 0. The total cost = 0 + 1 = 1. Notice that workers with index 1 and 2 are common in the first and last 3 workers.
- In the second hiring round we choose the worker from [2,4,1]. The lowest cost is 1 (index 2). The total cost = 1 + 1 = 2.
- In the third hiring round there are less than three candidates. We choose the worker from the remaining workers [2,4]. The lowest cost is 2 (index 0). The total cost = 2 + 2 = 4.
The total hiring cost is 4.

Constraints:

  • 1 <= costs.length <= 105
  • 1 <= costs[i] <= 105
  • 1 <= k, candidates <= costs.length

Solution:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
class Solution:
    def totalCost(self, costs: List[int], k: int, candidates: int) -> int:
        q1 = []
        q2 = []
        i, j = 0, len(costs) - 1

        ans = 0
        while k > 0:
            # construct two queues w/o common elements
            while len(q1) < candidates and i <= j:
                heapq.heappush(q1, costs[i])
                i += 1
            while len(q2) < candidates and i <= j:
                heapq.heappush(q2, costs[j])
                j -= 1

            # peek the smallest element from each qeue
            cost1 = q1[0] if q1 else float("inf")
            cost2 = q2[0] if q2 else float("inf")

            # select smaller cost from two queues
            if cost1 <= cost2:
                ans += cost1
                heapq.heappop(q1)
            else:
                ans += cost2
                heapq.heappop(q2)

            k -= 1

        return ans

[Leetcode 2542] Maximum Subsequence Score

You are given two 0-indexed integer arrays nums1 and nums2 of equal length n and a positive integer k. You must choose a subsequence of indices from nums1 of length k.

For chosen indices i0, i1, …, ik - 1, your score is defined as:

  • The sum of the selected elements from nums1 multiplied with the minimum of the selected elements from nums2.
  • It can defined simply as: (nums1[i0] + nums1[i1] +…+ nums1[ik - 1]) * min(nums2[i0] , nums2[i1], … ,nums2[ik - 1]).
    Return the maximum possible score.

A subsequence of indices of an array is a set that can be derived from the set {0, 1, …, n-1} by deleting some or no elements.

Example 1:

1
2
3
4
5
6
7
8
9
Input: nums1 = [1,3,3,2], nums2 = [2,1,3,4], k = 3
Output: 12
Explanation: 
The four possible subsequence scores are:
- We choose the indices 0, 1, and 2 with score = (1+3+3) * min(2,1,3) = 7.
- We choose the indices 0, 1, and 3 with score = (1+3+2) * min(2,1,4) = 6. 
- We choose the indices 0, 2, and 3 with score = (1+3+2) * min(2,3,4) = 12. 
- We choose the indices 1, 2, and 3 with score = (3+3+2) * min(1,3,4) = 8.
Therefore, we return the max score, which is 12.

Constraints:

  • n == nums1.length == nums2.length
  • 1 <= n <= 105
  • 0 <= nums1[i], nums2[j] <= 105
  • 1 <= k <= n

Solution:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
class Solution:
    def maxScore(self, nums1: List[int], nums2: List[int], k: int) -> int:
        preSum = 0
        minHeap = []
        res = 0

        # 1. construct pair list by zipping nums1 and nums2, 
        # 2. sort the list by 2nd element in descending order
        # 3. calculate presum with the first elements in zipped list
        # 4. if there is k elements in minheap, multiply presum with the second element in each pair.
        # 5. the second elelment in each pair is the minimum element as we iterate so far
        # 6. pop the first element from minheap 
        for n1, n2 in sorted(list(zip(nums1, nums2)), key = lambda n : n[1], reverse = True):
            preSum += n1
            heapq.heappush(minHeap, n1)

            if len(minHeap) == k:
                res = max(res, preSum * n2)
                preSum -= heapq.heappop(minHeap)

        return res