Advantage Shuffle LeetCode solution in Python

1
class Solution:
2
    def advantageCount(self, nums1: List[int], nums2: List[int]) -> List[int]:
3
        # sort by ascend order
4
        nums1.sort()
5

6
        # heapq in python is a minimal heap, so if we need to use a maximal heap, the element in the queue should be like: (-num, index)
7
        q = []
8
        for i, n in enumerate(nums2):
9
            heapq.heappush(q, (-n, i))
10

11
        n = len(nums1)
12
        l = 0
13
        r = n - 1
14
        ans = [0] * n
15

16
        while q:
17
            # pop out the maximum number from nums2
18
            num, i = heapq.heappop(q)
19
            num = -num
20

21
            # no number in nums1 > max number in nums2, so grap the minimum numer from nums1
22
            if nums1[r] <= num:
23
                ans[i] = nums1[l]
24
                l += 1
25
            else:
26
                ans[i] = nums1[r]
27
                r -= 1
28

29
        return ans