"""03-15-2022 LeetCode 1057. Campus Bikes""" from collections import defaultdict import heapq workers = [[0, 0], [2, 1]] bikes = [[1, 2], [3, 3]] #The following is the optimal solution given, using a heap and some clever ordering: # Still generates every worker/bike combination, but storing these in a heap makes # picking them out in order faster worker_to_bike_list = [] pq = [] for w, wl in enumerate(workers): curr_worker_pairs = [] for b, bl in enumerate(bikes): curr_worker_pairs.append(((abs(wl[0] - bl[0]) + abs(wl[1] - bl[1])), w, b)) # Sort the worker_to_bike_list for the current worker in reverse order curr_worker_pairs.sort(reverse=True) # Add the closest bike for this worker to the priority queue heapq.heappush(pq, curr_worker_pairs.pop()) # Store the remaining options for the current worker in worker_to_bike_list worker_to_bike_list.append(curr_worker_pairs) # Initialize all values to false, to signify no bikes have been taken bike_status = [False] * len(bikes) # Initialize all values to -1, to signify no worker has a bike worker_status = [-1] * len(workers) while pq: # Pop the worker-bike pair with smallest distance distance, worker, bike = heapq.heappop(pq) if not bike_status[bike]: # If the bike is free, assign the bike to the worker bike_status[bike] = True worker_status[worker] = bike else: # Otherwise, add the next closest bike for the current worker to the priority queue next_closest_bike = worker_to_bike_list[worker].pop() heapq.heappush(pq, next_closest_bike) print("return worker_status") # Damn! I was like 80% of the way there to using defaultdict to store w/b pairs by m_dist keys, # it seemed complicated and I didnt get the benefit at the time. The reason its better is that # you DONT need to sort: We are adding w/b pairs ALREADY in ascending order, only the m_dist changes # but by storing these in a dict, then iterating over the keys in the dict, we dont need to resort: # The m_dist is sorted naturally by iterating from min_m_dist through the rest of the dict, and each # val per key was added in already sorted. I did try saving m_dists to their own list and iterating # over just this, that way you dont bother checking w/b pairs for m_dists with no values. This served # to slow things down rather than speed things up: Presumably the w/b pairs are in a pretty small # tight order, and since m_dist is just an in dict check its very fast already. Seemed like making this # added set took more time than it saved. If m_dists were siginificantly more dispersed this may be a beter # choice overal. # Got the following to work # taken_workers = set() # taken_bikes = set() # d_w_b = defaultdict(list) # min_m_dist = 1998 # we know x, y <= 1000, 1000 so the largest distance is corner to corner 999 + 999 # wonb = [False]*len(workers) # for w in range(len(workers)): # for b in range(len(bikes)): # m_dist = abs(workers[w][0] - bikes[b][0]) + abs(workers[w][1] - bikes[b][1]) # min_m_dist = min(m_dist, min_m_dist) # d_w_b[m_dist].append((w, b)) # for i in range(min_m_dist, 1998): # for w, b in d_w_b[i]: # if w not in taken_workers and b not in taken_bikes: # wonb[w] = b # taken_workers.add(w) # taken_bikes.add(b) # if len(taken_workers) == len(workers): # break # return wonb # i += 1 # print(str(wonb)) # WOW! I can't believe I did this one. I had more intermediate steps that I thought needed to # be sorted each time, but it turns out .sort() automatically sorts multidimensional lists/tuples # correctly, including column by column # taken_workers = set() # taken_bikes = set() # d_w_b = [] # wonb = [False] * len(workers) # for i in range(len(workers)): # for j in range(len(bikes)): # d_w_b.append([abs(workers[i][0] - bikes[j][0]) + abs(workers[i][1] - bikes[j][1]),i, j]) # d_w_b.sort() # for d,w,b in d_w_b: # if w not in taken_workers and b not in taken_bikes: # wonb[w] = b # taken_workers.add(w) # taken_bikes.add(b) # if len(taken_workers) == len(workers): # break # print(str(wonb))