Problem statement

https://leetcode.com/problems/campus-bikes/

Solution

Let us have n workers and m bikes. Create all m*n pairs between workers and bikes: put them to heap and then extract one by one. Also we do lazy updates, that is we keep the set of used workers and bikes and pop from heap until we have not-used bike and person.

Complexity

It is O(nm*log(mn)) for time and O(mn) for space.

Code

class Solution:
    def assignBikes(self, workers, bikes):
        def dist(A, B):
            return abs(A[0] - B[0]) + abs(A[1] - B[1])

        n, m = len(workers), len(bikes)
        heap = []
        for i in range(n):
            for j in range(m):
                heap += [(dist(workers[i], bikes[j]), i, j)]

        heapify(heap)

        taken_W, taken_B = set(), set()
        ans = [0]*n
        for _ in range(n):
            while heap[0][1] in taken_W or heap[0][2] in taken_B:
                heappop(heap)
            dst, i, j = heappop(heap)
            taken_W.add(i)
            taken_B.add(j)
            ans[i] = j

        return ans

Remark

There is also O(mn) bucket sort solution: for each distnace put pairs of (i, j) ids of worker and bike to bucket of this distance. Then we iterate over buckets and iterate it until we meet the first not taken element. Note, that it can happen that we need to iterate full bucket if we are not lucky. I think it can be proven that there will be no more than O(m) element in each bucket, so overall time complexity is still O(mn), space is O(mn) as well.