squash_labs.py

class SearchTree:
    def __init__(self, cargos: List[dict], trucks: List[dict]):
        self.cargo_nodes, self.truck_nodes = self.build_nodes(cargos, trucks, Node)
        self.tree = KdTree(self.truck_nodes)
        self.delivery_distance = {
            item['product']: geodesic(
                (item['origin_lat'], item['origin_lng']),
                (item['destination_lat'], item['destination_lng'])
            ) for item in cargos
        }

    @staticmethod
    def build_nodes(cargos: List[dict], trucks: List[dict], node_class: Any) -> Tuple[List[Any], List[Any]]:
        cargo_nodes = [node_class((item['origin_lat'], item['origin_lng']), item['product']) for item in cargos]
        truck_nodes = [node_class((item['lat'], item['lng']), item['truck']) for item in trucks]

        return cargo_nodes, truck_nodes

    def truck_assignment(self) -> List[dict]:
        result = []
        for cargo in self.cargo_nodes:
            path, truck = self.tree.get_closest(cargo)
            result.append({
                "cargo": cargo.name,
                "truck": truck.name,
                "distance": round(path[truck].miles + self.delivery_distance[cargo.name].miles, 2)
            })

        return sorted(result, key=lambda x: x['distance'])