Problem statement

https://leetcode.com/problems/relative-sort-array/

Solution

Create dictionary, where for each value in B we keep its index. Then when we sort we use get function, which will return value if it is in dictionary and shifted by big value.

Complexity

It is O(n log n) for time and O(n) for space, where n = len(A).

Code

class Solution:
    def relativeSortArray(self, A, B):
        k = {b: i for i, b in enumerate(B)}
        return sorted(A, key=lambda a: k.get(a, 10**9 + a))