Untitled

import heapq

class CityMap:
    def __init__(self, roads, tracks, friends):
        self.location_index = []
        self.road_adj = []
        self.track_adj = []
        self.initialize_city(roads, tracks)
        self.friends = [(name, self.location_index.index(loc)) for name, loc in friends]

    def initialize_city(self, roads, tracks):
        # Initialize location mapping and adjacency lists
        location_set = set()
        for road in roads + tracks:
            location_set.update(road[:2])
        self.location_index = list(location_set)
        self.road_adj = [[] for _ in range(len(self.location_index))]
        self.track_adj = [[] for _ in range(len(self.location_index))]

        # Mapping locations to indices
        loc_to_index = {loc: idx for idx, loc in enumerate(self.location_index)}

        for u, v, w in roads:
            self.road_adj[loc_to_index[u]].append((loc_to_index[v], w))
            self.road_adj[loc_to_index[v]].append((loc_to_index[u], w))
        
        for u, v, w in tracks:
            self.track_adj[loc_to_index[u]].append((loc_to_index[v], w))

    def dijkstra(self, adj_list, start_index):
        dist = [float('inf')] * len(self.location_index)
        pred = [-1] * len(self.location_index)
        dist[start_index] = 0
        pq = [(0, start_index)]
        heapq.heapify(pq)

        while pq:
            current_dist, u = heapq.heappop(pq)
            if current_dist > dist[u]:
                continue

            for v, weight in adj_list[u]:
                if dist[v] > current_dist + weight:
                    dist[v] = current_dist + weight
                    pred[v] = u
                    heapq.heappush(pq, (dist[v], v))

        return dist, pred

    def plan(self, start, destination):
        start_index = self.location_index.index(start)
        destination_index = self.location_index.index(destination)
        road_distances, road_pred = self.dijkstra(self.road_adj, start_index)

        best_time = float('inf')
        best_route = []
        best_friend = None
        best_pickup_index = None

        for name, friend_loc in self.friends:
            track_distances, track_pred = self.dijkstra(self.track_adj, friend_loc)
            for i in range(len(self.location_index)):
                if road_distances[i] + track_distances[i] < best_time:
                    best_time = road_distances[i] + track_distances[i]
                    route_from_start = self.reconstruct_path(road_pred, start_index, i)
                    route_from_friend = self.reconstruct_path(track_pred, friend_loc, i)
                    best_route = route_from_start + route_from_friend[1:]
                    best_friend = name
                    best_pickup_index = i

        return best_time, best_route, best_friend, self.location_index[best_pickup_index]

    def reconstruct_path(self, pred, start, end):
        path = []
        current = end
        while current != start:
            path.append(self.location_index[current])
            current = pred[current]
        path.append(self.location_index[start])
        path.reverse()
        return path