[LeetCode]Advantage Shuffle@Python

Advantage Shuffle

Given two arrays A and B of equal size, the advantage of A with respect to B is the number of indices i for which A[i] > B[i].
Return any permutation of A that maximizes its advantage with respect to B.

Example

Input: A = [2,7,11,15], B = [1,10,4,11]
Output: [2,11,7,15]

Solution

class Solution:
    def advantageCount(self, A: List[int], B: List[int]) -> List[int]:
        A = sorted(A)
        take = collections.defaultdict(list)
        for b in sorted(B)[::-1]:
            if b<A[-1]:
                take[b].append(A.pop())
        return [(take[b] or A).pop() for b in B]